Skip to content

Commit 4940b0a

Browse files
authored
improvement(chat): coalesce replayed transcript reloads and narrow text-only message reads (#8977)
* improvement(chat): coalesce replayed transcript reloads and narrow text-only message reads * improvement(chat): bound a hung transcript reload and tighten text-block filtering - a chat detail reload in flight holds back later refresh requests for at most 30s; past it the next request starts its own reload and the stuck one is dropped when it settles - quoteChatTranscript filters blocks with isRecordLike - simplify the refresh coalescer's map init * fix(chat): cancel a stuck first transcript load before reloading, and skip malformed blocks - a reload past the hold cancels the stuck detail fetch first, since a query with no data joins its fetch in flight instead of replacing it - the text-block JSONPath runs silent, so a malformed block is skipped rather than failing the read - keep the non-array contentBlocks fixture out of the full-reader comparison, which cannot read it - shrink the unused-exports baseline now that collectChatFiles has an importer
1 parent 25da162 commit 4940b0a

15 files changed

Lines changed: 573 additions & 42 deletions

File tree

‎apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-registry/resource-invalidation.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import type { MothershipResourceType } from '@/lib/mothership/resources/types'
33
import { dashboardKeys } from '@/hooks/queries/dashboards'
44
import { deploymentKeys, invalidateDeploymentQueries } from '@/hooks/queries/deployments'
55
import { logKeys } from '@/hooks/queries/logs'
6-
import { mothershipChatKeys } from '@/hooks/queries/mothership-chats'
6+
import { mothershipChatKeys, refreshMothershipChatDetail } from '@/hooks/queries/mothership-chats'
77
import { folderKeys } from '@/hooks/queries/utils/folder-keys'
88
import { invalidateWorkflowLists } from '@/hooks/queries/utils/invalidate-workflow-lists'
99
import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
@@ -59,7 +59,7 @@ const RESOURCE_INVALIDATORS: Record<
5959
},
6060
task: (qc, wId, id) => {
6161
invalidate(qc, mothershipChatKeys.workspaceLists(wId))
62-
if (id) invalidate(qc, mothershipChatKeys.detail(id))
62+
if (id) refreshMothershipChatDetail(qc, id)
6363
},
6464
log: (qc, wId, id) => {
6565
invalidate(qc, [...logKeys.lists(), wId])

‎apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts‎

Lines changed: 136 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { createDeferred } from '@sim/testing/helpers/deferred'
12
import { QueryClient, QueryObserver } from '@tanstack/react-query'
23
import { beforeEach, describe, expect, it, vi } from 'vitest'
34
import { resourceScopeKey } from '@/lib/core/resource-scope'
@@ -26,7 +27,10 @@ vi.mock('@/hooks/queries/utils/workflow-cache', () => ({
2627

2728
import type { PersistedStreamEventEnvelope } from '@/lib/mothership/request/session/contract'
2829
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'
3034
import {
3135
makeStreamLoopDeps,
3236
ref,
@@ -197,6 +201,137 @@ describe('live resource effects', () => {
197201
)
198202
})
199203

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+
200335
describe('handleResourceEvent saved-view pins', () => {
201336
beforeEach(() => {
202337
useTableViewPinStore.getState().reset()

‎apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.ts‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,11 @@ import type {
2424
MothershipResource,
2525
MothershipResourceType,
2626
} from '@/app/workspace/[workspaceId]/home/types'
27-
import { type MothershipChatHistory, mothershipChatKeys } from '@/hooks/queries/mothership-chats'
27+
import {
28+
type MothershipChatHistory,
29+
mothershipChatKeys,
30+
refreshMothershipChatDetail,
31+
} from '@/hooks/queries/mothership-chats'
2832
import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
2933
import { removeWorkflowFromActiveCache } from '@/hooks/queries/utils/workflow-cache'
3034
import { useTableViewPinStore } from '@/stores/table/view-pin/store'
@@ -63,6 +67,18 @@ function applyResourceEffect(
6367
}
6468
}
6569

70+
/**
71+
* Recovery reloads the current panel state from the stored chat instead of
72+
* repeating historical focus commands. A batch replay reloads once at its
73+
* terminal flush; a reconnected live tail coalesces its reloads.
74+
*/
75+
function refreshRecoveredChatResources(ctx: StreamLoopContext): void {
76+
const chatId = ctx.deps.chatIdRef.current
77+
if (!chatId) return
78+
if (ctx.deps.options.deferFlushes) ctx.state.deferredChatDetailRefreshes.add(chatId)
79+
else refreshMothershipChatDetail(ctx.deps.queryClient, chatId)
80+
}
81+
6682
/**
6783
* Applies a streamed resource upsert/remove to the mothership resource list,
6884
* reconciling it with any in-flight or just-completed file-preview handoff so a
@@ -159,9 +175,7 @@ export function handleResourceEvent(ctx: StreamLoopContext, parsed: ResourceEven
159175
}
160176
}
161177
if (payload.replay || ctx.deps.options.deferFlushes) {
162-
const chatId = ctx.deps.chatIdRef.current
163-
if (chatId)
164-
void queryClient.invalidateQueries({ queryKey: mothershipChatKeys.detail(chatId) })
178+
refreshRecoveredChatResources(ctx)
165179
return
166180
}
167181
if (ctx.deps.citedSourcesEnabled) return
@@ -190,9 +204,7 @@ export function handleResourceEvent(ctx: StreamLoopContext, parsed: ResourceEven
190204
}
191205
if (payload.op === 'refresh') return
192206
if (payload.replay || ctx.deps.options.deferFlushes) {
193-
/** Recovery refetches current panel state instead of repeating historical focus commands. */
194-
const chatId = ctx.deps.chatIdRef.current
195-
if (chatId) void queryClient.invalidateQueries({ queryKey: mothershipChatKeys.detail(chatId) })
207+
refreshRecoveredChatResources(ctx)
196208
return
197209
}
198210
if (payload.op === 'clear_view') {

‎apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-context.ts‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@ import type {
2222
MothershipResource,
2323
MothershipResourceType,
2424
} from '@/app/workspace/[workspaceId]/home/types'
25-
import type { MothershipChatHistory } from '@/hooks/queries/mothership-chats'
25+
import {
26+
type MothershipChatHistory,
27+
refreshMothershipChatDetail,
28+
} from '@/hooks/queries/mothership-chats'
2629

2730
/** Minimum spacing between text-driven snapshot flushes (see flushText). */
2831
const MIN_TEXT_FLUSH_INTERVAL_MS = 50
@@ -83,6 +86,8 @@ export interface StreamLoopState {
8386
scheduledTextFlushFrame: number | null
8487
/** Trailing timer for the min-interval text-flush gate (see flushText). */
8588
scheduledTextFlushTimer: ReturnType<typeof setTimeout> | null
89+
/** Chats a batch replay recovered resources for, reloaded once by its terminal flush. */
90+
deferredChatDetailRefreshes: Set<string>
8691
}
8792

8893
export interface StreamEventScope {
@@ -198,7 +203,10 @@ export interface StreamLoopOps {
198203
isStale: () => boolean
199204
flush: () => void
200205
flushText: () => void
201-
/** Real flush that bypasses `deferFlushes` — the batch-replay terminal flush. */
206+
/**
207+
* Real flush that bypasses `deferFlushes` — the batch-replay terminal flush — then
208+
* reloads each chat the replay recovered resources for.
209+
*/
202210
forceFlush: () => void
203211
}
204212

@@ -234,6 +242,7 @@ export function createStreamLoopContext(deps: StreamLoopDeps): StreamLoopContext
234242
browserAgentRunIds: new Set(),
235243
scheduledTextFlushFrame: null,
236244
scheduledTextFlushTimer: null,
245+
deferredChatDetailRefreshes: new Set(),
237246
}
238247

239248
if (preserveState) {
@@ -371,7 +380,13 @@ export function createStreamLoopContext(deps: StreamLoopDeps): StreamLoopContext
371380
flushAndStamp()
372381
},
373382
flushText,
374-
forceFlush: flush,
383+
forceFlush: () => {
384+
flush()
385+
for (const chatId of state.deferredChatDetailRefreshes) {
386+
refreshMothershipChatDetail(deps.queryClient, chatId)
387+
}
388+
state.deferredChatDetailRefreshes.clear()
389+
},
375390
}
376391

377392
return { state, ops, deps }

‎apps/sim/hooks/queries/mothership-chats.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,53 @@ export const mothershipChatKeys = {
109109
detail: (chatId: string | undefined) => [...mothershipChatKeys.details(), chatId ?? ''] as const,
110110
}
111111

112+
/**
113+
* How long a chat's reload in flight holds back later requests. Past it a hung fetch no longer
114+
* blocks the chat: the next request starts its own reload, which replaces the stuck one.
115+
*/
116+
const CHAT_DETAIL_REFRESH_HOLD_MS = 30 * 1000
117+
118+
interface ChatDetailRefresh {
119+
startedAt: number
120+
requestedAgain: boolean
121+
}
122+
123+
/** Per client, the chats with a reload in flight. */
124+
const chatDetailRefreshes = new WeakMap<QueryClient, Map<string, ChatDetailRefresh>>()
125+
126+
/**
127+
* Invalidates a chat's detail so a mounted transcript reloads, coalescing a burst of
128+
* requests (a replayed or reconnected stream re-announces every resource it touched)
129+
* into one reload in flight plus one trailing reload. The trailing reload starts after
130+
* the last request, so the final read reflects every change requested meanwhile; a read
131+
* already in flight may predate a request and is never trusted to include it. A reload
132+
* in flight longer than {@link CHAT_DETAIL_REFRESH_HOLD_MS} stops holding requests back.
133+
*/
134+
export function refreshMothershipChatDetail(queryClient: QueryClient, chatId: string): void {
135+
const refreshes = chatDetailRefreshes.get(queryClient) ?? new Map<string, ChatDetailRefresh>()
136+
chatDetailRefreshes.set(queryClient, refreshes)
137+
const inFlight = refreshes.get(chatId)
138+
if (inFlight && Date.now() - inFlight.startedAt < CHAT_DETAIL_REFRESH_HOLD_MS) {
139+
inFlight.requestedAgain = true
140+
return
141+
}
142+
const refresh: ChatDetailRefresh = { startedAt: Date.now(), requestedAgain: false }
143+
refreshes.set(chatId, refresh)
144+
const queryKey = mothershipChatKeys.detail(chatId)
145+
// A query with no data yet joins its fetch in flight instead of replacing it, so a stuck
146+
// one is cancelled first.
147+
const reload = inFlight
148+
? queryClient
149+
.cancelQueries({ queryKey, exact: true })
150+
.then(() => queryClient.invalidateQueries({ queryKey }))
151+
: queryClient.invalidateQueries({ queryKey })
152+
void reload.finally(() => {
153+
if (refreshes.get(chatId) !== refresh) return
154+
refreshes.delete(chatId)
155+
if (refresh.requestedAgain) refreshMothershipChatDetail(queryClient, chatId)
156+
})
157+
}
158+
112159
/** Shared by the `useMothershipChats` hook and the workspace sidebar prefetch. */
113160
export const MOTHERSHIP_CHAT_LIST_STALE_TIME = 60 * 1000
114161
export const MOTHERSHIP_CHAT_HISTORY_STALE_TIME = 30 * 1000

0 commit comments

Comments
 (0)