fix(tui): batch streaming part deltas to keep TUI responsive during fast streams

packages/tui/src/context/sync.tsx

@@ -193,6 +193,45 @@ export const {
         .then((x) => (x.data ?? []).toSorted((a, b) => a.id.localeCompare(b.id)))
     }
 
+    // Fast-streaming models emit message.part.delta faster than the renderer
+    // keeps up; applying every chunk synchronously starves input handling and
+    // freezes scrolling (#36043). Deltas accumulate per part field and flush
+    // together on a short interval instead.
+    const pendingDeltas = new Map<string, { messageID: string; partID: string; field: string; delta: string }>()
+    let deltaTimer: ReturnType<typeof setTimeout> | undefined
+    // Apply buffered deltas to the store. With no argument, flush everything (timer
+    // tick, stream completion). With a part, flush only that part's buffered deltas —
+    // used before applying a full snapshot so accumulated text is committed rather
+    // than discarded.
+    const flushDeltas = (only?: { messageID: string; partID: string }) => {
+      if (only === undefined && deltaTimer) {
+        clearTimeout(deltaTimer)
+        deltaTimer = undefined
+      }
+      if (!pendingDeltas.size) return
+      const prefix = only ? `${only.messageID}|${only.partID}|` : undefined
+      batch(() => {
+        for (const [key, pending] of [...pendingDeltas.entries()]) {
+          if (prefix && !key.startsWith(prefix)) continue
+          pendingDeltas.delete(key)
+          const parts = store.part[pending.messageID]
+          if (!parts) continue
+          const result = search(parts, pending.partID, (p) => p.id)
+          if (!result.found) continue
+          setStore(
+            "part",
+            pending.messageID,
+            produce((draft) => {
+              const part = draft[result.index]
+              const field = pending.field as keyof typeof part
+              const existing = part[field] as string | undefined
+              ;(part[field] as string) = (existing ?? "") + pending.delta
+            }),
+          )
+        }
+      })
+    }
+
     event.subscribe((event, { directory, workspace }) => {
       switch (event.type) {
         case "server.instance.disposed":
@@ -334,11 +373,17 @@ export const {
         }
 
         case "session.status": {
+          // When a session goes idle the stream is done; flush any buffered deltas so
+          // a stream that ended on a delta (no trailing snapshot) isn't truncated.
+          if (event.properties.status.type === "idle") flushDeltas()
           setStore("session_status", event.properties.sessionID, event.properties.status)
           break
         }
 
         case "message.updated": {
+          // A completed assistant message ends its stream; flush buffered deltas so
+          // trailing streamed text is committed even without a final snapshot.
+          if (event.properties.info.role === "assistant" && event.properties.info.time?.completed) flushDeltas()
           if (event.properties.info.role === "user") {
             const optimisticID = optimisticBySession.get(event.properties.info.sessionID)
             if (optimisticID && optimisticID !== event.properties.info.id) {
@@ -400,6 +445,9 @@ export const {
           break
         }
         case "message.part.updated": {
+          // Commit this part's buffered deltas before the full snapshot reconciles,
+          // so accumulated streamed text is preserved rather than discarded.
+          flushDeltas({ messageID: event.properties.part.messageID, partID: event.properties.part.id })
           touchPart(event.properties.part.sessionID, event.properties.part.id)
           const parts = store.part[event.properties.part.messageID]
           if (!parts) {
@@ -427,16 +475,17 @@ export const {
           const result = search(parts, event.properties.partID, (p) => p.id)
           if (!result.found) break
           touchPart(event.properties.sessionID, event.properties.partID)
-          setStore(
-            "part",
-            event.properties.messageID,
-            produce((draft) => {
-              const part = draft[result.index]
-              const field = event.properties.field as keyof typeof part
-              const existing = part[field] as string | undefined
-              ;(part[field] as string) = (existing ?? "") + event.properties.delta
-            }),
-          )
+          const key = `${event.properties.messageID}|${event.properties.partID}|${event.properties.field}`
+          const pending = pendingDeltas.get(key)
+          if (pending) pending.delta += event.properties.delta
+          else
+            pendingDeltas.set(key, {
+              messageID: event.properties.messageID,
+              partID: event.properties.partID,
+              field: event.properties.field,
+              delta: event.properties.delta,
+            })
+          if (!deltaTimer) deltaTimer = setTimeout(flushDeltas, 40)
           break
         }
 

packages/tui/test/cli/cmd/tui/sync-live-hydration.test.tsx

@@ -390,3 +390,92 @@ test("hydration returning the real user message drops the optimistic duplicate",
     app.renderer.destroy()
   }
 })
+
+const streaming = { ...assistant, time: { created: 1 } }
+
+async function mountSimple(tmpPath: string) {
+  return mount((url) => {
+    if (url.pathname === `/session/${sessionID}`) return json(session)
+    if (url.pathname === `/session/${sessionID}/todo` || url.pathname === `/session/${sessionID}/diff`) return json([])
+    return undefined
+  }, tmpPath)
+}
+
+function partText(sync: Awaited<ReturnType<typeof mountSimple>>["sync"]) {
+  const part = sync.data.part[messageID]?.[0]
+  return part && "text" in part ? part.text : undefined
+}
+
+test("a completed assistant message does not lose a trailing buffered delta", async () => {
+  await using tmp = await tmpdir()
+  await Bun.write(`${tmp.path}/kv.json`, "{}")
+  const { app, emit, sync } = await mountSimple(tmp.path)
+  try {
+    emit(global({ id: "m0", type: "message.updated", properties: { sessionID, info: streaming } }))
+    emit(
+      global({
+        id: "p0",
+        type: "message.part.updated",
+        properties: { sessionID, time: 1, part: { id: partID, sessionID, messageID, type: "text", text: "Hi" } },
+      }),
+    )
+    await wait(() => sync.data.part[messageID]?.[0]?.type === "text")
+
+    // Stream a delta then complete the message. The completion flush commits the
+    // trailing delta rather than dropping it; the final text must include it.
+    emit(
+      global({
+        id: "d0",
+        type: "message.part.delta",
+        properties: { sessionID, messageID, partID, field: "text", delta: " there" },
+      }),
+    )
+    emit(global({ id: "m1", type: "message.updated", properties: { sessionID, info: assistant } }))
+    await wait(() => partText(sync) === "Hi there")
+    expect(partText(sync)).toBe("Hi there")
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
+test("a full snapshot after a buffered delta reconciles without losing text", async () => {
+  await using tmp = await tmpdir()
+  await Bun.write(`${tmp.path}/kv.json`, "{}")
+  const { app, emit, sync } = await mountSimple(tmp.path)
+  try {
+    emit(global({ id: "m0", type: "message.updated", properties: { sessionID, info: streaming } }))
+    emit(
+      global({
+        id: "p0",
+        type: "message.part.updated",
+        properties: { sessionID, time: 1, part: { id: partID, sessionID, messageID, type: "text", text: "Hi" } },
+      }),
+    )
+    await wait(() => sync.data.part[messageID]?.[0]?.type === "text")
+
+    // Delta then a monotonic snapshot carrying the accumulated text: the buffered
+    // delta is committed before the snapshot reconciles, so nothing is lost.
+    emit(
+      global({
+        id: "d0",
+        type: "message.part.delta",
+        properties: { sessionID, messageID, partID, field: "text", delta: " there" },
+      }),
+    )
+    emit(
+      global({
+        id: "p1",
+        type: "message.part.updated",
+        properties: {
+          sessionID,
+          time: 2,
+          part: { id: partID, sessionID, messageID, type: "text", text: "Hi there!" },
+        },
+      }),
+    )
+    await wait(() => partText(sync) === "Hi there!")
+    expect(partText(sync)).toBe("Hi there!")
+  } finally {
+    app.renderer.destroy()
+  }
+})