fix(parts-batcher): defer operations when target message is missing from cache (#262)

* fix(parts-batcher): defer operations when target message is missing from cache

* fix(parts-batcher): add TTL for deferred operations to prevent repeated invalidations

Author: chriswritescode-dev

frontend/src/hooks/useSSE.test.tsx

@@ -4,7 +4,8 @@ import type { ReactNode } from 'react'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { useSSE } from './useSSE'
 import { useSessionStatus } from '../stores/sessionStatusStore'
-import type { Part, MessageWithParts } from '@/api/types'
+import type { MessageWithParts } from '@/api/types'
+import { createTextPart } from '@/lib/partsBatcher'
 
 const mocks = vi.hoisted(() => ({
   getSessionStatuses: vi.fn(),
@@ -293,14 +294,14 @@ describe('useSSE', () => {
         ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo-a'],
         [{
           ...assistantMessage('session-1', 'message-1'),
-          parts: [textPart('session-1', 'message-1', 'part-1', 'A')],
+          parts: [createTextPart('session-1', 'message-1', 'part-1', 'A')],
         }],
       )
       queryClient.setQueryData(
         ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo-b'],
         [{
           ...assistantMessage('session-1', 'message-1'),
-          parts: [textPart('session-1', 'message-1', 'part-1', 'B')],
+          parts: [createTextPart('session-1', 'message-1', 'part-1', 'B')],
         }],
       )
 
@@ -377,7 +378,7 @@ describe('useSSE', () => {
         ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
         [{
           ...assistantMessage('session-1', 'message-1'),
-          parts: [textPart('session-1', 'message-1', 'part-1', '')],
+          parts: [createTextPart('session-1', 'message-1', 'part-1', '')],
         }],
       )
 
@@ -450,6 +451,3 @@ function assistantMessage(sessionID: string, messageID: string): MessageWithPart
   }
 }
 
-function textPart(sessionID: string, messageID: string, partID: string, text: string): Part {
-  return { id: partID, sessionID, messageID, type: 'text', text } as Part
-}

frontend/src/lib/partsBatcher.test.ts

@@ -1,7 +1,7 @@
 import { QueryClient } from '@tanstack/react-query'
-import { describe, it, expect, vi } from 'vitest'
-import { createPartsBatcher } from './partsBatcher'
-import type { Part, MessageWithParts } from '@/api/types'
+import { describe, it, expect, vi, afterEach } from 'vitest'
+import { createPartsBatcher, createTextPart, DEFERRED_OPERATION_TTL_MS } from './partsBatcher'
+import type { MessageWithParts } from '@/api/types'
 
 function assistantMessage(sessionID: string, messageID: string): MessageWithParts {
   return {
@@ -23,21 +23,21 @@ function assistantMessage(sessionID: string, messageID: string): MessageWithPart
   }
 }
 
-function textPart(sessionID: string, messageID: string, partID: string, text: string): Part {
-  return { id: partID, sessionID, messageID, type: 'text', text } as Part
-}
-
 function createManyCachedMessages(count: number, sessionID: string): MessageWithParts[] {
   const messages: MessageWithParts[] = []
   for (let i = 0; i < count; i++) {
     const msg = assistantMessage(sessionID, `msg-${i}`)
-    msg.parts = [textPart(sessionID, `msg-${i}`, `part-${i}`, `base text ${i}`)]
+    msg.parts = [createTextPart(sessionID, `msg-${i}`, `part-${i}`, `base text ${i}`)]
     messages.push(msg)
   }
   return messages
 }
 
 describe('createPartsBatcher', () => {
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
   it('invalidates when part deltas arrive before message cache exists and applies a later authoritative upsert', () => {
     const queryClient = new QueryClient()
     const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
@@ -59,7 +59,7 @@ describe('createPartsBatcher', () => {
       [assistantMessage('session-1', 'message-1')],
     )
 
-    batcher.queuePartUpdate('session-1', textPart('session-1', 'message-1', 'part-1', 'Hello world'), '/repo')
+    batcher.queuePartUpdate('session-1', createTextPart('session-1', 'message-1', 'part-1', 'Hello world'), '/repo')
     batcher.flush()
 
     const data = queryClient.getQueryData<MessageWithParts[]>([
@@ -82,7 +82,7 @@ describe('createPartsBatcher', () => {
       [assistantMessage('session-1', 'message-1')],
     )
 
-    batcher.queuePartUpdate('session-1', textPart('session-1', 'message-1', 'part-1', 'authoritative text'), '/repo')
+    batcher.queuePartUpdate('session-1', createTextPart('session-1', 'message-1', 'part-1', 'authoritative text'), '/repo')
     batcher.flush()
 
     const data = queryClient.getQueryData<MessageWithParts[]>([
@@ -100,7 +100,6 @@ describe('createPartsBatcher', () => {
 
   it('does not replay unapplied deltas onto refetched authoritative data', () => {
     const queryClient = new QueryClient()
-    const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
     const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
 
     queryClient.setQueryData(
@@ -111,21 +110,91 @@ describe('createPartsBatcher', () => {
     batcher.queuePartDelta('session-1', 'message-1', 'part-1', 'text', ' stale', '/repo')
     batcher.flush()
 
+    let data = queryClient.getQueryData<MessageWithParts[]>([
+      'opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo',
+    ])
+    expect(data![0].parts[0]).toHaveProperty('text', ' stale')
+
+    queryClient.setQueryData(
+      ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
+      [{ ...assistantMessage('session-1', 'message-1'), parts: [createTextPart('session-1', 'message-1', 'part-1', 'fresh')] }],
+    )
+
+    batcher.flush()
+
+    data = queryClient.getQueryData<MessageWithParts[]>([
+      'opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo',
+    ])
+    expect(data![0].parts[0]).toHaveProperty('text', 'fresh')
+  })
+
+  it('keeps text deltas pending until a later message update creates the assistant message', () => {
+    const queryClient = new QueryClient()
+    const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
+    const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
+
+    queryClient.setQueryData(
+      ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
+      [assistantMessage('session-1', 'message-old')],
+    )
+
+    batcher.queuePartDelta('session-1', 'message-new', 'part-1', 'text', 'streamed', '/repo')
+    batcher.flush()
+
     expect(invalidateSpy).toHaveBeenCalledWith({
       queryKey: ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
     })
 
     queryClient.setQueryData(
       ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
-      [{ ...assistantMessage('session-1', 'message-1'), parts: [textPart('session-1', 'message-1', 'part-1', 'fresh')] }],
+      [assistantMessage('session-1', 'message-old'), assistantMessage('session-1', 'message-new')],
+    )
+
+    batcher.flush({ sessionID: 'session-1', directory: '/repo' })
+
+    const data = queryClient.getQueryData<MessageWithParts[]>([
+      'opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo',
+    ])
+    expect(data![1].parts).toHaveLength(1)
+    expect(data![1].parts[0]).toMatchObject({
+      id: 'part-1',
+      sessionID: 'session-1',
+      messageID: 'message-new',
+      type: 'text',
+      text: 'streamed',
+    })
+  })
+
+  it('invalidates once while deferring, then drops the operation after the TTL elapses for a never-arriving message', () => {
+    vi.useFakeTimers()
+    const queryClient = new QueryClient()
+    const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
+    const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
+
+    queryClient.setQueryData(
+      ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
+      [assistantMessage('session-1', 'message-old')],
     )
 
+    batcher.queuePartDelta('session-1', 'message-missing', 'part-1', 'text', 'streamed', '/repo')
     batcher.flush()
+    expect(invalidateSpy).toHaveBeenCalledTimes(1)
+
+    batcher.flush()
+    batcher.flush()
+    expect(invalidateSpy).toHaveBeenCalledTimes(1)
+
+    vi.advanceTimersByTime(DEFERRED_OPERATION_TTL_MS + 1)
+    batcher.flush()
+    batcher.flush()
+    expect(invalidateSpy).toHaveBeenCalledTimes(1)
 
     const data = queryClient.getQueryData<MessageWithParts[]>([
       'opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo',
     ])
-    expect(data![0].parts[0]).toHaveProperty('text', 'fresh')
+    expect(data).toHaveLength(1)
+    expect(data![0].info.id).toBe('message-old')
+    expect(data![0].parts).toHaveLength(0)
   })
 
   it('applies deltas queued after an authoritative upsert in the same batch', () => {
@@ -137,7 +206,7 @@ describe('createPartsBatcher', () => {
       [assistantMessage('session-1', 'message-1')],
     )
 
-    batcher.queuePartUpdate('session-1', textPart('session-1', 'message-1', 'part-1', 'snapshot'), '/repo')
+    batcher.queuePartUpdate('session-1', createTextPart('session-1', 'message-1', 'part-1', 'snapshot'), '/repo')
     batcher.queuePartDelta('session-1', 'message-1', 'part-1', 'text', ' later', '/repo')
     batcher.flush()
 
@@ -202,8 +271,8 @@ describe('createPartsBatcher', () => {
       [{
         ...assistantMessage('session-1', 'message-1'),
         parts: [
-          textPart('session-1', 'message-1', 'part-1', 'first'),
-          textPart('session-1', 'message-1', 'part-2', 'second'),
+          createTextPart('session-1', 'message-1', 'part-1', 'first'),
+          createTextPart('session-1', 'message-1', 'part-2', 'second'),
         ],
       }],
     )
@@ -227,11 +296,11 @@ describe('createPartsBatcher', () => {
 
     queryClient.setQueryData(
       ['opencode', 'messages', 'http://localhost:5551', 'session-a', '/repo-a'],
-      [{ ...assistantMessage('session-a', 'msg-1'), parts: [textPart('session-a', 'msg-1', 'part-1', 'A1')] }],
+      [{ ...assistantMessage('session-a', 'msg-1'), parts: [createTextPart('session-a', 'msg-1', 'part-1', 'A1')] }],
     )
     queryClient.setQueryData(
       ['opencode', 'messages', 'http://localhost:5551', 'session-b', '/repo-b'],
-      [{ ...assistantMessage('session-b', 'msg-2'), parts: [textPart('session-b', 'msg-2', 'part-2', 'B1')] }],
+      [{ ...assistantMessage('session-b', 'msg-2'), parts: [createTextPart('session-b', 'msg-2', 'part-2', 'B1')] }],
     )
 
     batcher.queuePartDelta('session-a', 'msg-1', 'part-1', 'text', ' delta A', '/repo-a')
@@ -264,7 +333,7 @@ describe('createPartsBatcher', () => {
 
     queryClient.setQueryData(
       ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
-      [{ ...assistantMessage('session-1', 'msg-1'), parts: [textPart('session-1', 'msg-1', 'part-1', 'text')] }],
+      [{ ...assistantMessage('session-1', 'msg-1'), parts: [createTextPart('session-1', 'msg-1', 'part-1', 'text')] }],
     )
 
     batcher.queuePartDelta('session-1', 'msg-1', 'part-1', 'text', ' updated', '/repo')
@@ -286,7 +355,7 @@ describe('createPartsBatcher', () => {
     )
     queryClient.setQueryData(
       ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo-b'],
-      [{ ...assistantMessage('session-1', 'message-1'), parts: [textPart('session-1', 'message-1', 'part-1', 'B')] }],
+      [{ ...assistantMessage('session-1', 'message-1'), parts: [createTextPart('session-1', 'message-1', 'part-1', 'B')] }],
     )
 
     batcher.queuePartDelta('session-1', 'message-1', 'part-1', 'text', ' + chunk', '/repo-b')