From 413fbb8714ecbbd11d3f699665065fac8f0423a0 Mon Sep 17 00:00:00 2001 From: KorenKrita Date: Mon, 3 Aug 2026 09:25:45 +0800 Subject: [PATCH] fix(web): preserve queued edits across cancellation (#1324) --- .../AssistantChat/QueuedMessagesBar.test.tsx | 454 +++++++++++++++++- .../AssistantChat/QueuedMessagesBar.tsx | 231 +++++++-- web/src/components/SessionChat.tsx | 23 +- web/src/lib/locales/en.ts | 1 + web/src/lib/locales/zh-CN.ts | 1 + web/src/lib/queued-edit-recovery.test.ts | 103 ++++ web/src/lib/queued-edit-recovery.ts | 207 ++++++++ 7 files changed, 981 insertions(+), 39 deletions(-) create mode 100644 web/src/lib/queued-edit-recovery.test.ts create mode 100644 web/src/lib/queued-edit-recovery.ts diff --git a/web/src/components/AssistantChat/QueuedMessagesBar.test.tsx b/web/src/components/AssistantChat/QueuedMessagesBar.test.tsx index be981572..a84da2e1 100644 --- a/web/src/components/AssistantChat/QueuedMessagesBar.test.tsx +++ b/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( + + ) + return { + onEdit, + unmount: view.unmount, + rerender: (nextPendingSchedule: PendingSchedule | null, nextPendingScheduleRevision = currentPendingScheduleRevision) => { + currentPendingScheduleRevision = nextPendingScheduleRevision + view.rerender( + + ) + }, + } +} + +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((resolve, reject) => { + mocks.resolveCancel = resolve + mocks.rejectCancel = reject + })) +}) + +async function resolveCancel(result: DeferredCancelResult): Promise { + await act(async () => { + mocks.resolveCancel?.(result) + await Promise.resolve() + }) +} + +function installManualAnimationFrames() { + let nextHandle = 1 + const pending = new Map() + const history = new Map() + 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( + + ) + + 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 diff --git a/web/src/components/AssistantChat/QueuedMessagesBar.tsx b/web/src/components/AssistantChat/QueuedMessagesBar.tsx index 00504a75..384b2314 100644 --- a/web/src/components/AssistantChat/QueuedMessagesBar.tsx +++ b/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()) + // 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() + 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 diff --git a/web/src/components/SessionChat.tsx b/web/src/components/SessionChat.tsx index 256f8d6f..59a2d5d5 100644 --- a/web/src/components/SessionChat.tsx +++ b/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(null) + const [pendingScheduleRevision, setPendingScheduleRevision] = useState(0) + const updatePendingSchedule = useCallback((next: PendingSchedule | null) => { + setPendingSchedule(next) + setPendingScheduleRevision((revision) => revision + 1) + }, []) const pendingScheduleRef = useRef(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).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) { { // Restore the schedule so the clock button re-activates - setPendingSchedule(restored) + updatePendingSchedule(restored) }} /> @@ -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} diff --git a/web/src/lib/locales/en.ts b/web/src/lib/locales/en.ts index 9d51e7d6..f0970ab1 100644 --- a/web/src/lib/locales/en.ts +++ b/web/src/lib/locales/en.ts @@ -530,6 +530,7 @@ export default { 'composer.scheduleErrorTooFar': 'Maximum schedule time is 7 days.', 'queuedMessages.scheduledFor': 'Scheduled for {time}', 'queuedMessages.editAlreadyInvoked': "Message already sent — it can't be edited", + 'queuedMessages.editCurrentDraftKept': 'Queued message cancelled — current draft and schedule were kept.', // Scratchlist (per-session workbench, issue #11) 'scratchlist.title': 'Scratchlist', diff --git a/web/src/lib/locales/zh-CN.ts b/web/src/lib/locales/zh-CN.ts index 447e7057..2bf7eb7c 100644 --- a/web/src/lib/locales/zh-CN.ts +++ b/web/src/lib/locales/zh-CN.ts @@ -534,6 +534,7 @@ export default { 'composer.scheduleErrorTooFar': '最多只能定时 7 天。', 'queuedMessages.scheduledFor': '定时发送: {time}', 'queuedMessages.editAlreadyInvoked': '消息已发送，无法编辑', + 'queuedMessages.editCurrentDraftKept': '队列消息已取消，已保留当前草稿和定时设置。', // Scratchlist (per-session workbench, issue #11) 'scratchlist.title': '草稿夹', diff --git a/web/src/lib/queued-edit-recovery.test.ts b/web/src/lib/queued-edit-recovery.test.ts new file mode 100644 index 00000000..af9309c7 --- /dev/null +++ b/web/src/lib/queued-edit-recovery.test.ts @@ -0,0 +1,103 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const STORAGE_KEY = 'hapi:queued-edit-recovery' + +async function loadStore() { + return await import('./queued-edit-recovery') +} + +beforeEach(() => { + sessionStorage.clear() + vi.resetModules() +}) + +describe('queued-edit-recovery', () => { + it('drops invalid persisted schedules and rewrites the sanitized cache', async () => { + sessionStorage.setItem(STORAGE_KEY, JSON.stringify({ + badPreset: { + id: 'bad-preset', + text: 'text', + pendingSchedule: { type: 'preset', preset: '+99m' }, + composerTextAtEdit: '', + pendingScheduleAtEdit: null, + }, + badAbsolute: { + id: 'bad-absolute', + text: 'text', + pendingSchedule: { type: 'absolute', ms: Infinity }, + composerTextAtEdit: '', + pendingScheduleAtEdit: null, + }, + })) + + const store = await loadStore() + expect(store.getQueuedEditRecovery('badPreset')).toBeNull() + expect(store.getQueuedEditRecovery('badAbsolute')).toBeNull() + expect(sessionStorage.getItem(STORAGE_KEY)).toBe('{}') + }) + + it('notifies session listeners and returns isolated recovery copies', async () => { + const store = await loadStore() + const listener = vi.fn() + const unsubscribe = store.subscribeQueuedEditRecovery('session-1', listener) + + store.saveQueuedEditRecovery('session-1', { + text: 'queued edit', + pendingSchedule: { type: 'preset', preset: '+30m' }, + composerTextAtEdit: 'before edit', + pendingScheduleAtEdit: null, + }) + + expect(listener).toHaveBeenCalledTimes(1) + expect(store.isQueuedOperationPending('session-1')).toBe(true) + expect(store.beginQueuedOperation('session-1')).toBeNull() + const first = store.getQueuedEditRecovery('session-1')! + first.pendingSchedule = null + expect(store.getQueuedEditRecovery('session-1')?.pendingSchedule).toEqual({ type: 'preset', preset: '+30m' }) + + unsubscribe() + store.saveQueuedEditRecovery('session-1', { + text: 'latest edit', + pendingSchedule: null, + composerTextAtEdit: '', + pendingScheduleAtEdit: null, + }) + expect(listener).toHaveBeenCalledTimes(1) + }) + + it('keeps at most one pending operation per session and token-protects newer work', async () => { + const store = await loadStore() + const listener = vi.fn() + const unsubscribe = store.subscribeQueuedOperation('session-1', listener) + const first = store.beginQueuedOperation('session-1') + + expect(first).not.toBeNull() + expect(store.isQueuedOperationPending('session-1')).toBe(true) + expect(store.beginQueuedOperation('session-1')).toBeNull() + expect(listener).toHaveBeenCalledTimes(1) + + store.endQueuedOperation('session-1', first!) + const second = store.beginQueuedOperation('session-1')! + store.endQueuedOperation('session-1', first!) + expect(store.isQueuedOperationPending('session-1')).toBe(true) + store.endQueuedOperation('session-1', second) + expect(store.isQueuedOperationPending('session-1')).toBe(false) + expect(listener).toHaveBeenCalledTimes(4) + unsubscribe() + }) + + it('retains every unconsumed recovery beyond the former session cap', async () => { + const store = await loadStore() + for (let index = 0; index <= 50; index++) { + store.saveQueuedEditRecovery(`session-${index}`, { + text: String(index), + pendingSchedule: null, + composerTextAtEdit: '', + pendingScheduleAtEdit: null, + }) + } + + expect(store.getQueuedEditRecovery('session-0')?.text).toBe('0') + expect(store.getQueuedEditRecovery('session-50')?.text).toBe('50') + }) +}) diff --git a/web/src/lib/queued-edit-recovery.ts b/web/src/lib/queued-edit-recovery.ts new file mode 100644 index 00000000..78ad450d --- /dev/null +++ b/web/src/lib/queued-edit-recovery.ts @@ -0,0 +1,207 @@ +import type { PendingSchedule } from '@/components/AssistantChat/ScheduleTimePicker' + +const STORAGE_KEY = 'hapi:queued-edit-recovery' + +export type QueuedEditRecovery = { + id: string + text: string + pendingSchedule: PendingSchedule | null + composerTextAtEdit: string + pendingScheduleAtEdit: PendingSchedule | null +} +export type QueuedEditRecoveryInput = Omit + +type RecoveryMap = Record +export type QueuedOperationToken = symbol + +let cache: RecoveryMap | null = null +let recoverySequence = 0 +const listeners = new Map void>>() +const pendingOperationTokens = new Map() +const pendingOperationListeners = new Map void>>() + +const PRESETS = new Set['preset']>([ + '+5m', + '+30m', + '+1h', + '+4h', +]) + +function clonePendingSchedule(schedule: PendingSchedule | null): PendingSchedule | null { + if (schedule === null) return null + return schedule.type === 'preset' + ? { type: 'preset', preset: schedule.preset } + : { type: 'absolute', ms: schedule.ms } +} + +function isPendingSchedule(value: unknown): value is PendingSchedule { + if (!value || typeof value !== 'object') return false + const record = value as Record + return (record.type === 'preset' && typeof record.preset === 'string' && PRESETS.has(record.preset as never)) + || (record.type === 'absolute' && typeof record.ms === 'number' && Number.isFinite(record.ms)) +} + +function parsePendingSchedule(value: unknown): PendingSchedule | null | undefined { + if (value === null) return null + if (!isPendingSchedule(value)) return undefined + return clonePendingSchedule(value) +} + +function hydrate(): RecoveryMap { + if (cache) return cache + if (typeof window === 'undefined') { + cache = {} + return cache + } + try { + const raw = sessionStorage.getItem(STORAGE_KEY) + if (!raw) { + cache = {} + return cache + } + const parsed = JSON.parse(raw) as unknown + if (!parsed || typeof parsed !== 'object') { + cache = {} + return cache + } + const result: RecoveryMap = {} + let sanitized = false + for (const [sessionId, value] of Object.entries(parsed as Record)) { + if (!sessionId.trim() || !value || typeof value !== 'object') { + sanitized = true + continue + } + const record = value as Record + const pendingSchedule = parsePendingSchedule(record.pendingSchedule) + const pendingScheduleAtEdit = parsePendingSchedule(record.pendingScheduleAtEdit) + if ( + typeof record.id !== 'string' + || + typeof record.text !== 'string' + || typeof record.composerTextAtEdit !== 'string' + || pendingSchedule === undefined + || pendingScheduleAtEdit === undefined + ) { + sanitized = true + continue + } + result[sessionId] = { + id: record.id, + text: record.text, + pendingSchedule, + composerTextAtEdit: record.composerTextAtEdit, + pendingScheduleAtEdit, + } + } + cache = result + if (sanitized) { + try { + sessionStorage.setItem(STORAGE_KEY, JSON.stringify(result)) + } catch { + // Keep the in-memory sanitized cache even if persistence fails. + } + } + return cache + } catch { + cache = {} + return cache + } +} + +function persist(): void { + if (typeof window === 'undefined') return + try { + sessionStorage.setItem(STORAGE_KEY, JSON.stringify(hydrate())) + } catch { + // Recovery is best effort when sessionStorage is unavailable or full. + } +} + +function notify(sessionId: string): void { + for (const listener of listeners.get(sessionId) ?? []) { + listener() + } +} + +function notifyPendingOperation(sessionId: string): void { + for (const listener of pendingOperationListeners.get(sessionId) ?? []) { + listener() + } +} + +export function getQueuedEditRecovery(sessionId: string): QueuedEditRecovery | null { + const recovery = hydrate()[sessionId] + if (!recovery) return null + return { + id: recovery.id, + text: recovery.text, + pendingSchedule: clonePendingSchedule(recovery.pendingSchedule), + composerTextAtEdit: recovery.composerTextAtEdit, + pendingScheduleAtEdit: clonePendingSchedule(recovery.pendingScheduleAtEdit), + } +} + +export function saveQueuedEditRecovery(sessionId: string, recovery: QueuedEditRecoveryInput): void { + const recoveries = hydrate() + delete recoveries[sessionId] + recoveries[sessionId] = { + id: `${Date.now()}:${++recoverySequence}`, + text: recovery.text, + pendingSchedule: clonePendingSchedule(recovery.pendingSchedule), + composerTextAtEdit: recovery.composerTextAtEdit, + pendingScheduleAtEdit: clonePendingSchedule(recovery.pendingScheduleAtEdit), + } + persist() + notify(sessionId) + notifyPendingOperation(sessionId) +} + +export function clearQueuedEditRecovery(sessionId: string): void { + const recoveries = hydrate() + if (!recoveries[sessionId]) return + delete recoveries[sessionId] + persist() + notifyPendingOperation(sessionId) +} + +export function subscribeQueuedEditRecovery(sessionId: string, listener: () => void): () => void { + const sessionListeners = listeners.get(sessionId) ?? new Set<() => void>() + sessionListeners.add(listener) + listeners.set(sessionId, sessionListeners) + return () => { + sessionListeners.delete(listener) + if (sessionListeners.size === 0) listeners.delete(sessionId) + } +} + +/** + * Starts the single queued-message operation allowed for a session. + * The opaque token prevents an older completion from releasing a newer one. + */ +export function beginQueuedOperation(sessionId: string): QueuedOperationToken | null { + if (pendingOperationTokens.has(sessionId) || Boolean(hydrate()[sessionId])) return null + const token = Symbol(`queued-operation:${sessionId}`) + pendingOperationTokens.set(sessionId, token) + notifyPendingOperation(sessionId) + return token +} + +export function endQueuedOperation(sessionId: string, token: QueuedOperationToken): void { + if (pendingOperationTokens.get(sessionId) !== token) return + pendingOperationTokens.delete(sessionId) + notifyPendingOperation(sessionId) +} + +export function isQueuedOperationPending(sessionId: string): boolean { + return pendingOperationTokens.has(sessionId) || Boolean(hydrate()[sessionId]) +} + +export function subscribeQueuedOperation(sessionId: string, listener: () => void): () => void { + const sessionListeners = pendingOperationListeners.get(sessionId) ?? new Set<() => void>() + sessionListeners.add(listener) + pendingOperationListeners.set(sessionId, sessionListeners) + return () => { + sessionListeners.delete(listener) + if (sessionListeners.size === 0) pendingOperationListeners.delete(sessionId) + } +}