|
| 1 | +/** |
| 2 | + * @vitest-environment jsdom |
| 3 | + * |
| 4 | + * Regression tests for the mount-settling send loss: a send started on a |
| 5 | + * fresh chat surface used to be silently dropped when React ran the unmount |
| 6 | + * cleanup mid-flight (a Suspense hide/reveal cycles every effect shortly |
| 7 | + * after Home mounts), aborting the fetch before it dispatched. The fix routes |
| 8 | + * idle sends through the durable message queue and restores the queued entry |
| 9 | + * when the cleanup abort strikes before the server received the request. |
| 10 | + */ |
| 11 | +import { act, type ReactNode } from 'react' |
| 12 | +import { sleep } from '@sim/utils/helpers' |
| 13 | +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' |
| 14 | +import { createRoot, type Root } from 'react-dom/client' |
| 15 | +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' |
| 16 | + |
| 17 | +const { mockRequestJson, navigationMocks } = vi.hoisted(() => ({ |
| 18 | + mockRequestJson: vi.fn(), |
| 19 | + navigationMocks: { |
| 20 | + usePathname: vi.fn(() => '/workspace/ws-1/home'), |
| 21 | + useRouter: vi.fn(() => ({ push: vi.fn(), replace: vi.fn(), prefetch: vi.fn() })), |
| 22 | + useSearchParams: vi.fn(() => new URLSearchParams()), |
| 23 | + }, |
| 24 | +})) |
| 25 | + |
| 26 | +vi.mock('next/navigation', () => navigationMocks) |
| 27 | + |
| 28 | +vi.mock('@/lib/api/client/request', async (importOriginal) => ({ |
| 29 | + ...(await importOriginal<typeof import('@/lib/api/client/request')>()), |
| 30 | + requestJson: mockRequestJson, |
| 31 | +})) |
| 32 | + |
| 33 | +import { useChat } from '@/app/workspace/[workspaceId]/home/hooks/use-chat' |
| 34 | +import { useMothershipQueueStore } from '@/stores/mothership-queue/store' |
| 35 | + |
| 36 | +interface NetworkState { |
| 37 | + /** How the chat POST behaves for the next call. */ |
| 38 | + postBehavior: 'hang' | 'accept' |
| 39 | + postCalls: number |
| 40 | +} |
| 41 | + |
| 42 | +const state: NetworkState = { postBehavior: 'hang', postCalls: 0 } |
| 43 | + |
| 44 | +/** An SSE response whose stream ends immediately without a terminal event. */ |
| 45 | +function emptySseResponse(): Response { |
| 46 | + const stream = new ReadableStream<Uint8Array>({ |
| 47 | + start(controller) { |
| 48 | + controller.close() |
| 49 | + }, |
| 50 | + }) |
| 51 | + return new Response(stream, { status: 200, headers: { 'Content-Type': 'text/event-stream' } }) |
| 52 | +} |
| 53 | + |
| 54 | +function abortError(): Error { |
| 55 | + const error = new Error('Aborted') |
| 56 | + error.name = 'AbortError' |
| 57 | + return error |
| 58 | +} |
| 59 | + |
| 60 | +async function fetchStub(input: RequestInfo | URL, init?: RequestInit): Promise<Response> { |
| 61 | + const url = String(input instanceof Request ? input.url : input) |
| 62 | + |
| 63 | + if (url.includes('/api/mothership/chat') && init?.method === 'POST') { |
| 64 | + state.postCalls++ |
| 65 | + if (state.postBehavior === 'accept') return emptySseResponse() |
| 66 | + return new Promise<Response>((_, reject) => { |
| 67 | + const signal = init?.signal |
| 68 | + if (!signal) return |
| 69 | + if (signal.aborted) { |
| 70 | + reject(abortError()) |
| 71 | + return |
| 72 | + } |
| 73 | + signal.addEventListener('abort', () => reject(abortError()), { once: true }) |
| 74 | + }) |
| 75 | + } |
| 76 | + |
| 77 | + return new Response(JSON.stringify({ error: 'not found' }), { status: 404 }) |
| 78 | +} |
| 79 | + |
| 80 | +const mountedRoots: Root[] = [] |
| 81 | +let queryClient: QueryClient |
| 82 | + |
| 83 | +function renderUseChat(): { |
| 84 | + getResult: () => ReturnType<typeof useChat> |
| 85 | + unmount: () => void |
| 86 | +} { |
| 87 | + ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true |
| 88 | + queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }) |
| 89 | + const container = document.createElement('div') |
| 90 | + const root = createRoot(container) |
| 91 | + mountedRoots.push(root) |
| 92 | + let result: ReturnType<typeof useChat> | undefined |
| 93 | + |
| 94 | + function Probe() { |
| 95 | + result = useChat('ws-1', undefined) |
| 96 | + return null |
| 97 | + } |
| 98 | + |
| 99 | + act(() => { |
| 100 | + root.render( |
| 101 | + <QueryClientProvider client={queryClient}>{(<Probe />) as ReactNode}</QueryClientProvider> |
| 102 | + ) |
| 103 | + }) |
| 104 | + |
| 105 | + return { |
| 106 | + getResult: () => { |
| 107 | + if (result === undefined) throw new Error('Hook result is not ready') |
| 108 | + return result |
| 109 | + }, |
| 110 | + unmount: () => act(() => root.unmount()), |
| 111 | + } |
| 112 | +} |
| 113 | + |
| 114 | +/** Every queued message across all chat keys, flattened. */ |
| 115 | +function allQueuedMessages() { |
| 116 | + return Object.values(useMothershipQueueStore.getState().queues).flat() |
| 117 | +} |
| 118 | + |
| 119 | +async function waitFor(predicate: () => boolean, budgetMs = 2000): Promise<void> { |
| 120 | + const deadline = Date.now() + budgetMs |
| 121 | + while (!predicate()) { |
| 122 | + if (Date.now() > deadline) throw new Error('waitFor timed out') |
| 123 | + await act(async () => { |
| 124 | + await sleep(10) |
| 125 | + }) |
| 126 | + } |
| 127 | +} |
| 128 | + |
| 129 | +describe('useChat mount-settling send recovery', () => { |
| 130 | + beforeEach(() => { |
| 131 | + vi.stubGlobal('fetch', fetchStub) |
| 132 | + state.postBehavior = 'hang' |
| 133 | + state.postCalls = 0 |
| 134 | + mockRequestJson.mockResolvedValue({ chats: [] }) |
| 135 | + useMothershipQueueStore.setState({ queues: {}, editing: {} }) |
| 136 | + window.sessionStorage.clear() |
| 137 | + }) |
| 138 | + |
| 139 | + afterEach(() => { |
| 140 | + for (const root of mountedRoots.splice(0)) { |
| 141 | + act(() => root.unmount()) |
| 142 | + } |
| 143 | + queryClient?.clear() |
| 144 | + vi.unstubAllGlobals() |
| 145 | + vi.clearAllMocks() |
| 146 | + }) |
| 147 | + |
| 148 | + it('restores a send the unmount cleanup aborted before the server received it', async () => { |
| 149 | + const { getResult, unmount } = renderUseChat() |
| 150 | + |
| 151 | + await act(async () => { |
| 152 | + void getResult().sendMessage('hello from the palette') |
| 153 | + }) |
| 154 | + await waitFor(() => state.postCalls === 1) |
| 155 | + |
| 156 | + // The dispatch claimed the queue head when the optimistic send applied. |
| 157 | + expect(allQueuedMessages()).toHaveLength(0) |
| 158 | + |
| 159 | + // The cleanup abort (same code path a Suspense hide/reveal runs during |
| 160 | + // mount-settling) fires while the POST is still awaiting the server. |
| 161 | + unmount() |
| 162 | + await waitFor(() => allQueuedMessages().length === 1) |
| 163 | + |
| 164 | + expect(allQueuedMessages()[0].content).toBe('hello from the palette') |
| 165 | + }) |
| 166 | + |
| 167 | + it('does not re-queue a send the server already received', async () => { |
| 168 | + state.postBehavior = 'accept' |
| 169 | + const { getResult, unmount } = renderUseChat() |
| 170 | + |
| 171 | + await act(async () => { |
| 172 | + void getResult().sendMessage('already accepted') |
| 173 | + }) |
| 174 | + await waitFor(() => state.postCalls === 1) |
| 175 | + |
| 176 | + unmount() |
| 177 | + await act(async () => { |
| 178 | + await sleep(50) |
| 179 | + }) |
| 180 | + |
| 181 | + expect(allQueuedMessages()).toHaveLength(0) |
| 182 | + }) |
| 183 | +}) |
0 commit comments