refactor(sse): batch parts directory-aware and serialize backend writes

Author: chriswritescode-dev

frontend/src/components/message/MessageThread.tsx

@@ -135,8 +135,7 @@ const findLastMessageByRole = (
 
 interface MessageRowProps {
   msgWithParts: MessageWithParts
-  index: number
-  messages: MessageWithParts[]
+  nextAssistantMessage: MessageWithParts | undefined
   pendingAssistantId: string | undefined
   lastUserMessageId: string | undefined
   isSessionBusy: boolean
@@ -157,8 +156,7 @@ interface MessageRowProps {
 
 const MessageRow = memo(function MessageRow({
   msgWithParts,
-  index,
-  messages,
+  nextAssistantMessage,
   pendingAssistantId,
   lastUserMessageId,
   isSessionBusy,
@@ -183,7 +181,6 @@ const MessageRow = memo(function MessageRow({
   const isLastUserMessage = msg.role === 'user' && msg.id === lastUserMessageId
   const messageTextContent = getMessageTextContent(parts)
 
-  const nextAssistantMessage = messages.slice(index + 1).find(m => m.info.role === 'assistant')
   const nextAssistantMsg = nextAssistantMessage?.info
   const isUserBeforeAssistant = msg.role === 'user' && nextAssistantMessage
   const canEditUserMessage = isLastUserMessage && isUserBeforeAssistant && !isSessionBusy
@@ -352,6 +349,20 @@ export const MessageThread = memo(function MessageThread({
     return findLastMessageByRole(messages, 'user')
   }, [messages])
 
+  const nextAssistantByMessageId = useMemo(() => {
+    const map = new Map<string, MessageWithParts | undefined>()
+    if (!messages) return map
+    let next: MessageWithParts | undefined
+    for (let i = messages.length - 1; i >= 0; i--) {
+      const msg = messages[i]
+      map.set(msg.info.id, next)
+      if (msg.info.role === 'assistant') {
+        next = msg
+      }
+    }
+    return map
+  }, [messages])
+
   const isSessionBusy = !!pendingAssistantId || isSessionInRetry(sessionStatus)
   const setSessionTodos = useSessionTodos((state) => state.setTodos)
 
@@ -415,12 +426,11 @@ export const MessageThread = memo(function MessageThread({
 
   return (
     <div className="flex flex-col space-y-2 p-2 overflow-x-hidden">
-      {messages.map((msgWithParts, index) => (
+      {messages.map((msgWithParts) => (
         <MessageRow
           key={msgWithParts.info.id}
           msgWithParts={msgWithParts}
-          index={index}
-          messages={messages}
+          nextAssistantMessage={nextAssistantByMessageId.get(msgWithParts.info.id)}
           pendingAssistantId={pendingAssistantId}
           lastUserMessageId={lastUserMessageId}
           isSessionBusy={isSessionBusy}

frontend/src/hooks/useAutoScroll.test.tsx

@@ -131,6 +131,9 @@ describe('useAutoScroll', () => {
         contentVersion: newMessages.length,
         onScrollStateChange,
       })
+    })
+
+    act(() => {
       vi.advanceTimersByTime(100)
     })
 
@@ -162,6 +165,7 @@ describe('useAutoScroll', () => {
         contentVersion: messages.length + 1,
         onScrollStateChange,
       })
+      vi.advanceTimersByTime(100)
     })
 
     expect(containerHarness.getScrollTop()).toBe(containerHarness.div.scrollHeight - containerHarness.div.clientHeight)
@@ -231,6 +235,29 @@ describe('useAutoScroll', () => {
     expect(containerHarness.getScrollTop()).toBe(userPosition)
   })
 
+  it('cancels pending bottom scroll when user wheel-scrolls up before pending frames complete', () => {
+    const messages = [createMessage('1', 'user'), createMessage('2', 'assistant')]
+    const { renderResult, containerHarness } = setupHook(messages)
+
+    act(() => {
+      renderResult.result.current.scrollToBottom()
+    })
+
+    const userPosition = 150
+    act(() => {
+      containerHarness.div.dispatchEvent(
+        new WheelEvent('wheel', {
+          deltaY: -50,
+          bubbles: true,
+        })
+      )
+      containerHarness.setScrollTop(userPosition)
+      vi.runOnlyPendingTimers()
+    })
+
+    expect(containerHarness.getScrollTop()).toBe(userPosition)
+  })
+
   it('does not show scroll button on tiny upward drag from bottom', () => {
     const messages = [createMessage('1', 'user')]
     const { containerHarness, onScrollStateChange } = setupHook(messages)
@@ -431,6 +458,10 @@ describe('useAutoScroll', () => {
       })
     })
 
+    act(() => {
+      vi.advanceTimersByTime(100)
+    })
+
     expect(containerHarness.getScrollTop()).toBe(containerHarness.div.scrollHeight - containerHarness.div.clientHeight)
   })
 
@@ -459,6 +490,10 @@ describe('useAutoScroll', () => {
       })
     })
 
+    act(() => {
+      vi.advanceTimersByTime(100)
+    })
+
     expect(containerHarness.getScrollTop()).toBe(containerHarness.div.scrollHeight - containerHarness.div.clientHeight)
   })
 })

frontend/src/hooks/useAutoScroll.ts

@@ -46,17 +46,13 @@ export function useAutoScroll({
     isScrollButtonVisibleRef.current = false
     const scrollRequestId = scrollRequestIdRef.current + 1
     scrollRequestIdRef.current = scrollRequestId
-    const scroll = () => {
-      if (!containerRef?.current) return
-      containerRef.current.scrollTop = containerRef.current.scrollHeight
-    }
 
-    scroll()
     let frameCount = 0
     const scrollAfterLayout = () => {
       if (scrollRequestIdRef.current !== scrollRequestId) return
+      if (!containerRef?.current) return
+      containerRef.current.scrollTop = containerRef.current.scrollHeight
       frameCount += 1
-      scroll()
       if (frameCount < SCROLL_TO_BOTTOM_FRAME_COUNT) {
         requestAnimationFrame(scrollAfterLayout)
       }

frontend/src/lib/partsBatcher.test.ts

@@ -27,6 +27,16 @@ function textPart(sessionID: string, messageID: string, partID: string, text: st
   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}`)]
+    messages.push(msg)
+  }
+  return messages
+}
+
 describe('createPartsBatcher', () => {
   it('invalidates when part deltas arrive before message cache exists and applies a later authoritative upsert', () => {
     const queryClient = new QueryClient()
@@ -138,6 +148,79 @@ describe('createPartsBatcher', () => {
     expect(data![0].parts[0]).toHaveProperty('text', 'snapshot later')
   })
 
+  it('applies many queued part deltas with one cache write and no invalidation storm', () => {
+    const queryClient = new QueryClient()
+    const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
+
+    const sessionID = 'session-1'
+    const directory = '/repo'
+    const messageCount = 1000
+    const deltaCount = 500
+
+    const messages = createManyCachedMessages(messageCount, sessionID)
+    queryClient.setQueryData(
+      ['opencode', 'messages', 'http://localhost:5551', sessionID, directory],
+      messages,
+    )
+
+    const setQueryDataSpy = vi.spyOn(queryClient, 'setQueryData')
+    const invalidateSpy = vi.spyOn(queryClient, 'invalidateQueries')
+
+    for (let i = 0; i < deltaCount; i++) {
+      batcher.queuePartDelta(sessionID, `msg-${i}`, `part-${i}`, 'text', ` delta ${i}`, directory)
+    }
+
+    batcher.flush()
+
+    expect(invalidateSpy).not.toHaveBeenCalled()
+
+    const data = queryClient.getQueryData<MessageWithParts[]>([
+      'opencode', 'messages', 'http://localhost:5551', sessionID, directory,
+    ])
+    expect(data).toHaveLength(messageCount)
+
+    for (let i = 0; i < deltaCount; i++) {
+      expect(data![i].parts[0]).toHaveProperty('text', `base text ${i} delta ${i}`)
+    }
+
+    for (let i = deltaCount; i < messageCount; i++) {
+      expect(data![i].parts[0]).toHaveProperty('text', `base text ${i}`)
+    }
+
+    const setQueryDataCalls = setQueryDataSpy.mock.calls.filter(
+      ([key]) => JSON.stringify(key).includes('opencode'),
+    )
+    expect(setQueryDataCalls.length).toBe(1)
+  })
+
+  it('does not apply same-batch deltas for removed parts to shifted parts', () => {
+    const queryClient = new QueryClient()
+    const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')
+
+    queryClient.setQueryData(
+      ['opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo'],
+      [{
+        ...assistantMessage('session-1', 'message-1'),
+        parts: [
+          textPart('session-1', 'message-1', 'part-1', 'first'),
+          textPart('session-1', 'message-1', 'part-2', 'second'),
+        ],
+      }],
+    )
+
+    batcher.queuePartRemoval('session-1', 'message-1', 'part-1', '/repo')
+    batcher.queuePartDelta('session-1', 'message-1', 'part-1', 'text', ' stale', '/repo')
+    batcher.flush()
+
+    const data = queryClient.getQueryData<MessageWithParts[]>([
+      'opencode', 'messages', 'http://localhost:5551', 'session-1', '/repo',
+    ])
+
+    expect(data![0].parts).toHaveLength(1)
+    expect(data![0].parts[0]).toHaveProperty('id', 'part-2')
+    expect(data![0].parts[0]).toHaveProperty('text', 'second')
+  })
+
   it('applies deltas to the directory they were queued for', () => {
     const queryClient = new QueryClient()
     const batcher = createPartsBatcher(queryClient, 'http://localhost:5551')