|
| 1 | +import { createDeferred } from '@sim/testing/helpers/deferred' |
1 | 2 | import { QueryClient, QueryObserver } from '@tanstack/react-query' |
2 | 3 | import { beforeEach, describe, expect, it, vi } from 'vitest' |
3 | 4 | import { resourceScopeKey } from '@/lib/core/resource-scope' |
@@ -26,7 +27,10 @@ vi.mock('@/hooks/queries/utils/workflow-cache', () => ({ |
26 | 27 |
|
27 | 28 | import type { PersistedStreamEventEnvelope } from '@/lib/mothership/request/session/contract' |
28 | 29 | import { handleResourceEvent } from '@/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event' |
29 | | -import type { StreamLoopContext } from '@/app/workspace/[workspaceId]/home/hooks/stream/stream-context' |
| 30 | +import { |
| 31 | + createStreamLoopContext, |
| 32 | + type StreamLoopContext, |
| 33 | +} from '@/app/workspace/[workspaceId]/home/hooks/stream/stream-context' |
30 | 34 | import { |
31 | 35 | makeStreamLoopDeps, |
32 | 36 | ref, |
@@ -197,6 +201,137 @@ describe('live resource effects', () => { |
197 | 201 | ) |
198 | 202 | }) |
199 | 203 |
|
| 204 | +describe('recovered resource events', () => { |
| 205 | + const detailKey = mothershipChatKeys.detail('chat-1') |
| 206 | + const tables = ['tbl-1', 'tbl-2', 'tbl-3'].map( |
| 207 | + (id): MothershipResource => ({ type: 'table', id, title: id }) |
| 208 | + ) |
| 209 | + |
| 210 | + function replayed(resource: MothershipResource, seq: number): ResourceEvent { |
| 211 | + return { |
| 212 | + ...removeEvent('file', resource.id), |
| 213 | + seq, |
| 214 | + payload: { op: 'upsert', replay: true, resource }, |
| 215 | + } as ResourceEvent |
| 216 | + } |
| 217 | + |
| 218 | + /** A mounted chat whose every read returns the resources stored when the read began. */ |
| 219 | + function mountedChat(queryClient: QueryClient, { cached = true, hangFirstRead = false } = {}) { |
| 220 | + const server = { stored: [] as MothershipResource[] } |
| 221 | + const reads = vi.fn( |
| 222 | + async (): Promise<MothershipChatHistory> => ({ |
| 223 | + id: 'chat-1', |
| 224 | + title: null, |
| 225 | + messages: [], |
| 226 | + activeStreamId: 'stream-1', |
| 227 | + resources: [...server.stored], |
| 228 | + }) |
| 229 | + ) |
| 230 | + if (hangFirstRead) reads.mockImplementationOnce(() => new Promise(() => {})) |
| 231 | + if (cached) { |
| 232 | + queryClient.setQueryData(detailKey, { |
| 233 | + id: 'chat-1', |
| 234 | + title: null, |
| 235 | + messages: [], |
| 236 | + activeStreamId: 'stream-1', |
| 237 | + resources: [], |
| 238 | + } satisfies MothershipChatHistory) |
| 239 | + } |
| 240 | + const unsubscribe = new QueryObserver(queryClient, { |
| 241 | + queryKey: detailKey, |
| 242 | + queryFn: reads, |
| 243 | + staleTime: Number.POSITIVE_INFINITY, |
| 244 | + }).subscribe(() => {}) |
| 245 | + return { server, reads, unsubscribe } |
| 246 | + } |
| 247 | + |
| 248 | + async function settledResources(queryClient: QueryClient) { |
| 249 | + await vi.waitFor(() => expect(queryClient.isFetching({ queryKey: detailKey })).toBe(0)) |
| 250 | + return queryClient.getQueryData<MothershipChatHistory>(detailKey)?.resources |
| 251 | + } |
| 252 | + |
| 253 | + it('reloads the transcript once for a whole batch replay, after its terminal flush', async () => { |
| 254 | + const queryClient = new QueryClient() |
| 255 | + const { server, reads, unsubscribe } = mountedChat(queryClient) |
| 256 | + const ctx = createStreamLoopContext( |
| 257 | + makeStreamLoopDeps({ |
| 258 | + queryClient, |
| 259 | + chatIdRef: ref<string | undefined>('chat-1'), |
| 260 | + options: { deferFlushes: true }, |
| 261 | + }) |
| 262 | + ) |
| 263 | + server.stored = tables |
| 264 | + |
| 265 | + tables.forEach((table, index) => handleResourceEvent(ctx, replayed(table, index + 1))) |
| 266 | + expect(reads).not.toHaveBeenCalled() |
| 267 | + ctx.ops.forceFlush() |
| 268 | + |
| 269 | + expect(await settledResources(queryClient)).toEqual(tables) |
| 270 | + expect(reads).toHaveBeenCalledOnce() |
| 271 | + unsubscribe() |
| 272 | + }) |
| 273 | + |
| 274 | + it('coalesces a reconnected tail burst into one read plus one trailing read that sees the last change', async () => { |
| 275 | + const queryClient = new QueryClient() |
| 276 | + const { server, reads, unsubscribe } = mountedChat(queryClient) |
| 277 | + const deps = makeStreamLoopDeps({ queryClient, chatIdRef: ref<string | undefined>('chat-1') }) |
| 278 | + const ctx = createStreamLoopContext(deps) |
| 279 | + const firstReadReleased = createDeferred<void>() |
| 280 | + reads.mockImplementationOnce(async () => { |
| 281 | + const resources = [...server.stored] |
| 282 | + await firstReadReleased.promise |
| 283 | + return { id: 'chat-1', title: null, messages: [], activeStreamId: 'stream-1', resources } |
| 284 | + }) |
| 285 | + |
| 286 | + server.stored = tables.slice(0, 1) |
| 287 | + handleResourceEvent(ctx, replayed(tables[0], 1)) |
| 288 | + server.stored = tables |
| 289 | + handleResourceEvent(ctx, replayed(tables[1], 2)) |
| 290 | + handleResourceEvent(ctx, replayed(tables[2], 3)) |
| 291 | + firstReadReleased.resolve() |
| 292 | + |
| 293 | + await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(2)) |
| 294 | + expect(await settledResources(queryClient)).toEqual(tables) |
| 295 | + expect(reads).toHaveBeenCalledTimes(2) |
| 296 | + unsubscribe() |
| 297 | + }) |
| 298 | + |
| 299 | + it.each([ |
| 300 | + { name: 'a loaded transcript', cached: true }, |
| 301 | + { name: 'a transcript still on its first load', cached: false }, |
| 302 | + ])( |
| 303 | + 'stops holding reloads of $name behind one that never settles once 30 seconds pass', |
| 304 | + async ({ cached }) => { |
| 305 | + const queryClient = new QueryClient() |
| 306 | + const { server, reads, unsubscribe } = mountedChat(queryClient, { |
| 307 | + cached, |
| 308 | + hangFirstRead: true, |
| 309 | + }) |
| 310 | + const ctx = createStreamLoopContext( |
| 311 | + makeStreamLoopDeps({ queryClient, chatIdRef: ref<string | undefined>('chat-1') }) |
| 312 | + ) |
| 313 | + vi.useFakeTimers({ toFake: ['Date'] }) |
| 314 | + try { |
| 315 | + handleResourceEvent(ctx, replayed(tables[0], 1)) |
| 316 | + await vi.waitFor(() => expect(reads).toHaveBeenCalledOnce()) |
| 317 | + handleResourceEvent(ctx, replayed(tables[1], 2)) |
| 318 | + expect(reads).toHaveBeenCalledOnce() |
| 319 | + |
| 320 | + server.stored = tables |
| 321 | + vi.setSystemTime(Date.now() + 30_000) |
| 322 | + handleResourceEvent(ctx, replayed(tables[2], 3)) |
| 323 | + |
| 324 | + await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(2)) |
| 325 | + expect(await settledResources(queryClient)).toEqual(tables) |
| 326 | + expect(reads).toHaveBeenCalledTimes(2) |
| 327 | + } finally { |
| 328 | + vi.useRealTimers() |
| 329 | + unsubscribe() |
| 330 | + } |
| 331 | + } |
| 332 | + ) |
| 333 | +}) |
| 334 | + |
200 | 335 | describe('handleResourceEvent saved-view pins', () => { |
201 | 336 | beforeEach(() => { |
202 | 337 | useTableViewPinStore.getState().reset() |
|
0 commit comments