Skip to content

Commit bf51d58

Browse files
chriswritescode-devchriswritescode-dev
committed
fix(parts-batcher): add TTL for deferred operations to prevent repeated invalidations
1 parent 9c1b0aa commit bf51d58

3 files changed

Lines changed: 110 additions & 51 deletions

File tree

‎frontend/src/hooks/useSSE.test.tsx‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,8 @@ import type { ReactNode } from 'react'
44
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
55
import { useSSE } from './useSSE'
66
import { useSessionStatus } from '../stores/sessionStatusStore'
7-
import type { Part, MessageWithParts } from '@/api/types'
7+
import type { MessageWithParts } from '@/api/types'
8+
import { createTextPart } from '@/lib/partsBatcher'
89

910
const mocks = vi.hoisted(() => ({
1011
getSessionStatuses: vi.fn(),
@@ -293,14 +294,14 @@ describe('useSSE', () => {
293294
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo-a'],
294295
[{
295296
...assistantMessage('session-1', 'message-1'),
296-
parts: [textPart('session-1', 'message-1', 'part-1', 'A')],
297+
parts: [createTextPart('session-1', 'message-1', 'part-1', 'A')],
297298
}],
298299
)
299300
queryClient.setQueryData(
300301
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo-b'],
301302
[{
302303
...assistantMessage('session-1', 'message-1'),
303-
parts: [textPart('session-1', 'message-1', 'part-1', 'B')],
304+
parts: [createTextPart('session-1', 'message-1', 'part-1', 'B')],
304305
}],
305306
)
306307

@@ -377,7 +378,7 @@ describe('useSSE', () => {
377378
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
378379
[{
379380
...assistantMessage('session-1', 'message-1'),
380-
parts: [textPart('session-1', 'message-1', 'part-1', '')],
381+
parts: [createTextPart('session-1', 'message-1', 'part-1', '')],
381382
}],
382383
)
383384

@@ -450,6 +451,3 @@ function assistantMessage(sessionID: string, messageID: string): MessageWithPart
450451
}
451452
}
452453

453-
function textPart(sessionID: string, messageID: string, partID: string, text: string): Part {
454-
return { id: partID, sessionID, messageID, type: 'text', text } as Part
455-
}

‎frontend/src/lib/partsBatcher.test.ts‎

Lines changed: 50 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { QueryClient } from '@tanstack/react-query'
2-
import { describe, it, expect, vi } from 'vitest'
3-
import { createPartsBatcher } from './partsBatcher'
4-
import type { Part, MessageWithParts } from '@/api/types'
2+
import { describe, it, expect, vi, afterEach } from 'vitest'
3+
import { createPartsBatcher, createTextPart, DEFERRED_OPERATION_TTL_MS } from './partsBatcher'
4+
import type { MessageWithParts } from '@/api/types'
55

66
function assistantMessage(sessionID: string, messageID: string): MessageWithParts {
77
return {
@@ -23,21 +23,21 @@ function assistantMessage(sessionID: string, messageID: string): MessageWithPart
2323
}
2424
}
2525

26-
function textPart(sessionID: string, messageID: string, partID: string, text: string): Part {
27-
return { id: partID, sessionID, messageID, type: 'text', text } as Part
28-
}
29-
3026
function createManyCachedMessages(count: number, sessionID: string): MessageWithParts[] {
3127
const messages: MessageWithParts[] = []
3228
for (let i = 0; i < count; i++) {
3329
const msg = assistantMessage(sessionID, `msg-${i}`)
34-
msg.parts = [textPart(sessionID, `msg-${i}`, `part-${i}`, `base text ${i}`)]
30+
msg.parts = [createTextPart(sessionID, `msg-${i}`, `part-${i}`, `base text ${i}`)]
3531
messages.push(msg)
3632
}
3733
return messages
3834
}
3935

4036
describe('createPartsBatcher', () => {
37+
afterEach(() => {
38+
vi.useRealTimers()
39+
})
40+
4141
it('invalidates when part deltas arrive before message cache exists and applies a later authoritative upsert', () => {
4242
const queryClient = new QueryClient()
4343
const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
@@ -59,7 +59,7 @@ describe('createPartsBatcher', () => {
5959
[assistantMessage('session-1', 'message-1')],
6060
)
6161

62-
batcher.queuePartUpdate('session-1', textPart('session-1', 'message-1', 'part-1', 'Hello world'), '/repo')
62+
batcher.queuePartUpdate('session-1', createTextPart('session-1', 'message-1', 'part-1', 'Hello world'), '/repo')
6363
batcher.flush()
6464

6565
const data = queryClient.getQueryData<MessageWithParts[]>([
@@ -82,7 +82,7 @@ describe('createPartsBatcher', () => {
8282
[assistantMessage('session-1', 'message-1')],
8383
)
8484

85-
batcher.queuePartUpdate('session-1', textPart('session-1', 'message-1', 'part-1', 'authoritative text'), '/repo')
85+
batcher.queuePartUpdate('session-1', createTextPart('session-1', 'message-1', 'part-1', 'authoritative text'), '/repo')
8686
batcher.flush()
8787

8888
const data = queryClient.getQueryData<MessageWithParts[]>([
@@ -117,7 +117,7 @@ describe('createPartsBatcher', () => {
117117

118118
queryClient.setQueryData(
119119
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
120-
[{ ...assistantMessage('session-1', 'message-1'), parts: [textPart('session-1', 'message-1', 'part-1', 'fresh')] }],
120+
[{ ...assistantMessage('session-1', 'message-1'), parts: [createTextPart('session-1', 'message-1', 'part-1', 'fresh')] }],
121121
)
122122

123123
batcher.flush()
@@ -165,6 +165,38 @@ describe('createPartsBatcher', () => {
165165
})
166166
})
167167

168+
it('invalidates once while deferring, then drops the operation after the TTL elapses for a never-arriving message', () => {
169+
vi.useFakeTimers()
170+
const queryClient = new QueryClient()
171+
const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
172+
const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
173+
174+
queryClient.setQueryData(
175+
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
176+
[assistantMessage('session-1', 'message-old')],
177+
)
178+
179+
batcher.queuePartDelta('session-1', 'message-missing', 'part-1', 'text', 'streamed', '/repo')
180+
batcher.flush()
181+
expect(invalidateSpy).toHaveBeenCalledTimes(1)
182+
183+
batcher.flush()
184+
batcher.flush()
185+
expect(invalidateSpy).toHaveBeenCalledTimes(1)
186+
187+
vi.advanceTimersByTime(DEFERRED_OPERATION_TTL_MS + 1)
188+
batcher.flush()
189+
batcher.flush()
190+
expect(invalidateSpy).toHaveBeenCalledTimes(1)
191+
192+
const data = queryClient.getQueryData<MessageWithParts[]>([
193+
'opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo',
194+
])
195+
expect(data).toHaveLength(1)
196+
expect(data![0].info.id).toBe('message-old')
197+
expect(data![0].parts).toHaveLength(0)
198+
})
199+
168200
it('applies deltas queued after an authoritative upsert in the same batch', () => {
169201
const queryClient = new QueryClient()
170202
const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
@@ -174,7 +206,7 @@ describe('createPartsBatcher', () => {
174206
[assistantMessage('session-1', 'message-1')],
175207
)
176208

177-
batcher.queuePartUpdate('session-1', textPart('session-1', 'message-1', 'part-1', 'snapshot'), '/repo')
209+
batcher.queuePartUpdate('session-1', createTextPart('session-1', 'message-1', 'part-1', 'snapshot'), '/repo')
178210
batcher.queuePartDelta('session-1', 'message-1', 'part-1', 'text', ' later', '/repo')
179211
batcher.flush()
180212

@@ -239,8 +271,8 @@ describe('createPartsBatcher', () => {
239271
[{
240272
...assistantMessage('session-1', 'message-1'),
241273
parts: [
242-
textPart('session-1', 'message-1', 'part-1', 'first'),
243-
textPart('session-1', 'message-1', 'part-2', 'second'),
274+
createTextPart('session-1', 'message-1', 'part-1', 'first'),
275+
createTextPart('session-1', 'message-1', 'part-2', 'second'),
244276
],
245277
}],
246278
)
@@ -264,11 +296,11 @@ describe('createPartsBatcher', () => {
264296

265297
queryClient.setQueryData(
266298
['opencode', 'messages', 'http://localhost:5551', 'session-a', '/repo-a'],
267-
[{ ...assistantMessage('session-a', 'msg-1'), parts: [textPart('session-a', 'msg-1', 'part-1', 'A1')] }],
299+
[{ ...assistantMessage('session-a', 'msg-1'), parts: [createTextPart('session-a', 'msg-1', 'part-1', 'A1')] }],
268300
)
269301
queryClient.setQueryData(
270302
['opencode', 'messages', 'http://localhost:5551', 'session-b', '/repo-b'],
271-
[{ ...assistantMessage('session-b', 'msg-2'), parts: [textPart('session-b', 'msg-2', 'part-2', 'B1')] }],
303+
[{ ...assistantMessage('session-b', 'msg-2'), parts: [createTextPart('session-b', 'msg-2', 'part-2', 'B1')] }],
272304
)
273305

274306
batcher.queuePartDelta('session-a', 'msg-1', 'part-1', 'text', ' delta A', '/repo-a')
@@ -301,7 +333,7 @@ describe('createPartsBatcher', () => {
301333

302334
queryClient.setQueryData(
303335
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
304-
[{ ...assistantMessage('session-1', 'msg-1'), parts: [textPart('session-1', 'msg-1', 'part-1', 'text')] }],
336+
[{ ...assistantMessage('session-1', 'msg-1'), parts: [createTextPart('session-1', 'msg-1', 'part-1', 'text')] }],
305337
)
306338

307339
batcher.queuePartDelta('session-1', 'msg-1', 'part-1', 'text', ' updated', '/repo')
@@ -323,7 +355,7 @@ describe('createPartsBatcher', () => {
323355
)
324356
queryClient.setQueryData(
325357
['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo-b'],
326-
[{ ...assistantMessage('session-1', 'message-1'), parts: [textPart('session-1', 'message-1', 'part-1', 'B')] }],
358+
[{ ...assistantMessage('session-1', 'message-1'), parts: [createTextPart('session-1', 'message-1', 'part-1', 'B')] }],
327359
)
328360

329361
batcher.queuePartDelta('session-1', 'message-1', 'part-1', 'text', ' + chunk', '/repo-b')

‎frontend/src/lib/partsBatcher.ts‎

Lines changed: 55 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,11 @@ type OperationGroup = {
1919
sessionID: string
2020
directory?: string
2121
operations: PartOperation[]
22+
firstDeferredAt?: number
2223
}
2324

25+
export const DEFERRED_OPERATION_TTL_MS = 3000
26+
2427
function groupKey(sessionID: string, directory?: string): string {
2528
return `${directory ?? ''}\0${sessionID}`
2629
}
@@ -30,10 +33,32 @@ function deferOperation(operation: PartOperation): PartOperation | undefined {
3033
return { ...operation, deferred: true }
3134
}
3235

33-
function createTextPart(sessionID: string, messageID: string, partID: string, text: string): Part {
36+
function isGroupExpired(group: OperationGroup, now: number): boolean {
37+
return group.firstDeferredAt !== undefined && now - group.firstDeferredAt >= DEFERRED_OPERATION_TTL_MS
38+
}
39+
40+
function stampDeferred(group: OperationGroup, now: number): void {
41+
if (group.firstDeferredAt === undefined) group.firstDeferredAt = now
42+
}
43+
44+
export function createTextPart(sessionID: string, messageID: string, partID: string, text: string): Part {
3445
return { id: partID, sessionID, messageID, type: 'text', text } as Part
3546
}
3647

48+
function appendPart(
49+
updatedData: MessageWithParts[],
50+
msgIdx: number,
51+
msg: MessageWithParts,
52+
part: Part,
53+
pIdx: Map<string, number>,
54+
supersededPartIDs: Set<string>,
55+
): void {
56+
const nextParts = [...msg.parts, part]
57+
updatedData[msgIdx] = { ...msg, parts: nextParts }
58+
pIdx.set(part.id, nextParts.length - 1)
59+
supersededPartIDs.add(part.id)
60+
}
61+
3762
export function createPartsBatcher(
3863
queryClient: QueryClient,
3964
opcodeUrl: string,
@@ -52,8 +77,8 @@ export function createPartsBatcher(
5277
const flush = (target?: { sessionID?: string; directory?: string }) => {
5378
if (pendingOperations.size === 0) return
5479

80+
const now = Date.now()
5581
const groupsToDelete: string[] = []
56-
const invalidatedGroupKeys = new Set<string>()
5782

5883
for (const [key, group] of pendingOperations.entries()) {
5984
if (target) {
@@ -71,14 +96,16 @@ export function createPartsBatcher(
7196
const currentData = queryClient.getQueryData<MessageWithParts[]>(queryKey)
7297

7398
if (!currentData) {
74-
if (!invalidatedGroupKeys.has(key)) {
75-
invalidatedGroupKeys.add(key)
99+
if (group.firstDeferredAt === undefined) {
76100
queryClient.invalidateQueries({ queryKey })
77101
}
78-
group.operations = group.operations
102+
const deferred = group.operations
79103
.map(deferOperation)
80104
.filter((operation): operation is PartOperation => Boolean(operation))
81-
if (group.operations.length === 0) {
105+
if (deferred.length > 0 && !isGroupExpired(group, now)) {
106+
stampDeferred(group, now)
107+
group.operations = deferred
108+
} else {
82109
groupsToDelete.push(key)
83110
}
84111
continue
@@ -128,16 +155,14 @@ export function createPartsBatcher(
128155
supersededPartIDs.add(operation.part.id)
129156
continue
130157
}
131-
let nextParts: Part[]
132158
if (existingPartIdx !== undefined) {
133-
nextParts = [...msg.parts]
159+
const nextParts = [...msg.parts]
134160
nextParts[existingPartIdx] = operation.part
161+
updatedData[msgIdx] = { ...msg, parts: nextParts }
162+
supersededPartIDs.add(operation.part.id)
135163
} else {
136-
nextParts = [...msg.parts, operation.part]
137-
pIdx.set(operation.part.id, nextParts.length - 1)
164+
appendPart(updatedData, msgIdx, msg, operation.part, pIdx, supersededPartIDs)
138165
}
139-
updatedData[msgIdx] = { ...msg, parts: nextParts }
140-
supersededPartIDs.add(operation.part.id)
141166
continue
142167
}
143168

@@ -183,10 +208,14 @@ export function createPartsBatcher(
183208
const pIdxResult = pIdx.get(operation.partID)
184209
if (pIdxResult === undefined) {
185210
if (operation.field === 'text') {
186-
const nextParts = [...msg.parts, createTextPart(sessionID, operation.messageID, operation.partID, operation.delta)]
187-
updatedData[msgIdx] = { ...msg, parts: nextParts }
188-
pIdx.set(operation.partID, nextParts.length - 1)
189-
supersededPartIDs.add(operation.partID)
211+
appendPart(
212+
updatedData,
213+
msgIdx,
214+
msg,
215+
createTextPart(sessionID, operation.messageID, operation.partID, operation.delta),
216+
pIdx,
217+
supersededPartIDs,
218+
)
190219
} else {
191220
const deferred = deferOperation(operation)
192221
if (deferred) unapplied.push(deferred)
@@ -224,17 +253,17 @@ export function createPartsBatcher(
224253
return true
225254
})
226255

227-
const shouldInvalidate = retainedUnapplied.length > 0 || unapplied.some((op) => op.type === 'remove')
256+
const needsInvalidate = retainedUnapplied.length > 0 || unapplied.some((op) => op.type === 'remove')
257+
const willRetain = retainedUnapplied.length > 0 && !isGroupExpired(group, now)
228258

229-
if (shouldInvalidate) {
230-
if (!invalidatedGroupKeys.has(key)) {
231-
invalidatedGroupKeys.add(key)
232-
queryClient.invalidateQueries({ queryKey })
233-
}
234-
if (retainedUnapplied.length > 0) {
235-
group.operations = retainedUnapplied
236-
continue
237-
}
259+
if (needsInvalidate && group.firstDeferredAt === undefined) {
260+
queryClient.invalidateQueries({ queryKey })
261+
}
262+
263+
if (willRetain) {
264+
stampDeferred(group, now)
265+
group.operations = retainedUnapplied
266+
continue
238267
}
239268

240269
groupsToDelete.push(key)

0 commit comments

Comments
 (0)