fix(tui): coalesce message.part.delta store writes

packages/tui/src/context/delta-buffer.test.ts

@@ -0,0 +1,111 @@
+import { describe, expect, test } from "bun:test"
+import { createDeltaBuffer, type PendingDelta } from "./delta-buffer"
+
+function makeBuffer(history: PendingDelta[] = []) {
+  const calls: string[] = []
+  const queue: Array<() => void> = []
+  const buffer = createDeltaBuffer(
+    (flush) => {
+      calls.push("schedule")
+      queue.push(flush)
+    },
+    (item) => history.push(item),
+  )
+  const tick = () => queue.shift()?.()
+  return { buffer, queue, calls, tick }
+}
+
+const base = { sessionID: "s", messageID: "m" } as const
+
+describe("createDeltaBuffer", () => {
+  test("coalesces consecutive deltas for the same part into one application", () => {
+    const applied: PendingDelta[] = []
+    const { buffer, tick } = makeBuffer(applied)
+
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "a" })
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "b" })
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "c" })
+
+    expect(applied).toEqual([])
+    tick()
+    expect(applied).toEqual([{ kind: "reasoning", ...base, partID: "r1", delta: "abc" }])
+    for (const item of applied) expect(item).not.toBe(undefined)
+  })
+
+  test("keeps distinct parts and kinds separate within one flush", () => {
+    const applied: PendingDelta[] = []
+    const { buffer, tick } = makeBuffer(applied)
+
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "a" })
+    buffer.push({ kind: "text", ...base, partID: "t1", delta: "x" })
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "b" })
+    buffer.push({ kind: "tool", ...base, partID: "t1", delta: "y" })
+
+    tick()
+    expect(applied).toEqual([
+      { kind: "reasoning", ...base, partID: "r1", delta: "ab" },
+      { kind: "text", ...base, partID: "t1", delta: "x" },
+      { kind: "tool", ...base, partID: "t1", delta: "y" },
+    ])
+  })
+
+  test("drains pending deltas immediately without waiting for the schedule", () => {
+    const applied: PendingDelta[] = []
+    const { buffer, tick } = makeBuffer(applied)
+
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "a" })
+    buffer.drain()
+    expect(applied).toEqual([{ kind: "reasoning", ...base, partID: "r1", delta: "a" }])
+
+    tick()
+    expect(applied).toEqual([{ kind: "reasoning", ...base, partID: "r1", delta: "a" }])
+  })
+
+  test("a stale scheduled flush after drain is a no-op", () => {
+    const applied: PendingDelta[] = []
+    const { buffer, tick } = makeBuffer(applied)
+
+    buffer.push({ kind: "text", ...base, partID: "t1", delta: "a" })
+    buffer.drain()
+    buffer.push({ kind: "text", ...base, partID: "t1", delta: "b" })
+
+    tick()
+    expect(applied).toEqual([
+      { kind: "text", ...base, partID: "t1", delta: "a" },
+      { kind: "text", ...base, partID: "t1", delta: "b" },
+    ])
+  })
+
+  test("schedules at most one flush while deltas keep arriving", () => {
+    const { buffer, calls } = makeBuffer()
+
+    for (let i = 0; i < 100; i++) {
+      buffer.push({ kind: "text", ...base, partID: "t1", delta: String(i) })
+    }
+    expect(calls.length).toBe(1)
+  })
+
+  test("drop removes matching pending deltas without applying them", () => {
+    const applied: PendingDelta[] = []
+    const { buffer, tick } = makeBuffer(applied)
+
+    buffer.push({ kind: "text", ...base, partID: "t1", delta: "a" })
+    buffer.push({ kind: "text", ...base, partID: "t2", delta: "b" })
+    buffer.drop((item) => item.partID === "t1")
+
+    tick()
+    expect(applied).toEqual([{ kind: "text", ...base, partID: "t2", delta: "b" }])
+  })
+
+  test("drop for an entire message clears all of its parts", () => {
+    const applied: PendingDelta[] = []
+    const { buffer, tick } = makeBuffer(applied)
+
+    buffer.push({ kind: "text", ...base, partID: "t1", delta: "a" })
+    buffer.push({ kind: "reasoning", ...base, partID: "r1", delta: "x" })
+    buffer.drop((item) => item.messageID === "m")
+
+    tick()
+    expect(applied).toEqual([])
+  })
+})

packages/tui/src/context/delta-buffer.ts

@@ -0,0 +1,53 @@
+export type PendingDelta = {
+  kind: string
+  sessionID: string
+  messageID: string
+  partID: string
+  delta: string
+}
+
+/**
+ * Coalesces per-part stream deltas so a fast token stream rewrites shared store
+ * state at most once per scheduled flush instead of once per event. Reasoning
+ * blocks can stream thousands of deltas; applying each one immediately keeps
+ * re-rendering the accumulating text, which stalls the UI while expanded.
+ */
+export function createDeltaBuffer(schedule: (flush: () => void) => void, apply: (item: PendingDelta) => void) {
+  let pending = new Map<string, PendingDelta>()
+  let scheduled = false
+
+  const flush = () => {
+    scheduled = false
+    if (pending.size === 0) return
+    const items = [...pending.values()]
+    pending = new Map()
+    for (const item of items) apply(item)
+  }
+
+  return {
+    push(input: PendingDelta) {
+      const key = `${input.kind}:${input.sessionID}:${input.messageID}:${input.partID}`
+      const existing = pending.get(key)
+      if (existing) existing.delta += input.delta
+      else pending.set(key, { ...input })
+      if (!scheduled) {
+        scheduled = true
+        schedule(flush)
+      }
+    },
+    // Apply anything pending immediately so terminating events that replace the
+    // full text (text/reasoning/tool-input ended) observe the final deltas.
+    drain() {
+      flush()
+    },
+    // Discard pending deltas matched by `predicate` without applying them. Used
+    // when an authoritative full-part event (message.part.updated/removed,
+    // message.removed) supersedes streamed deltas so they are never applied to
+    // the already-replaced text.
+    drop(predicate: (item: PendingDelta) => boolean) {
+      for (const [key, item] of pending) {
+        if (predicate(item)) pending.delete(key)
+      }
+    },
+  }
+}

packages/tui/src/context/sync.tsx

@@ -28,7 +28,8 @@ import { useTuiStartup } from "./runtime"
 import { createSimpleContext } from "./helper"
 import { useExit } from "./exit"
 import { useArgs } from "./args"
-import { batch, onMount } from "solid-js"
+import { batch, onCleanup, onMount } from "solid-js"
+import { createDeltaBuffer } from "./delta-buffer"
 import path from "path"
 import { useKV } from "./kv"
 import { usePermission } from "./permission"
@@ -157,6 +158,34 @@ export const {
       hydratingSessions.get(sessionID)?.parts.add(partID)
     }
 
+    // Streaming text/reasoning deltas are coalesced before being applied to the
+    // store. Applying each message.part.delta immediately re-renders the whole
+    // growing part on every token, which is O(n²) in the part's length and can
+    // freeze the TUI at high stream rates. Buffer the appends and flush on a
+    // short timer, bounding renders to ~30fps regardless of token rate.
+    // Authoritative full-part events (message.part.updated / message.part.removed
+    // / message.removed) drop the buffer so superseded deltas are never applied
+    // to the already-replaced text.
+    const DELTA_COALESCE_MS = 32
+    const deltaBuffer = createDeltaBuffer(
+      (flush) => {
+        setTimeout(flush, DELTA_COALESCE_MS)
+      },
+      (item) => {
+        setStore(
+          "part",
+          item.messageID,
+          produce((draft) => {
+            const result = search(draft, item.partID, (part) => part.id)
+            if (!result.found) return
+            const part = draft[result.index]
+            if (!("text" in part)) return
+            part.text = (part.text ?? "") + item.delta
+          }),
+        )
+      },
+    )
+
     function sessionListQuery(): { scope?: "project"; path?: string } {
       if (!kv.get("session_directory_filter_enabled", true)) return { scope: "project" }
       if (!project.data.instance.path.worktree || !project.data.instance.path.directory) return { scope: "project" }
@@ -360,6 +389,7 @@ export const {
         }
         case "message.removed": {
           touchMessage(event.properties.sessionID, event.properties.messageID)
+          deltaBuffer.drop((item) => item.messageID === event.properties.messageID)
           const messages = store.message[event.properties.sessionID]
           const index = messages.findIndex((message) => message.id === event.properties.messageID)
           if (index !== -1) {
@@ -375,6 +405,10 @@ export const {
         }
         case "message.part.updated": {
           touchPart(event.properties.part.sessionID, event.properties.part.id)
+          deltaBuffer.drop(
+            (item) =>
+              item.messageID === event.properties.part.messageID && item.partID === event.properties.part.id,
+          )
           const parts = store.part[event.properties.part.messageID]
           if (!parts) {
             setStore("part", event.properties.part.messageID, [event.properties.part])
@@ -401,21 +435,21 @@ export const {
           const result = search(parts, event.properties.partID, (part) => part.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
-            }),
-          )
+          deltaBuffer.push({
+            kind: event.properties.field,
+            sessionID: event.properties.sessionID,
+            messageID: event.properties.messageID,
+            partID: event.properties.partID,
+            delta: event.properties.delta,
+          })
           break
         }
 
         case "message.part.removed": {
           touchPart(event.properties.sessionID, event.properties.partID)
+          deltaBuffer.drop(
+            (item) => item.messageID === event.properties.messageID && item.partID === event.properties.partID,
+          )
           const parts = store.part[event.properties.messageID]
           const result = search(parts, event.properties.partID, (part) => part.id)
           if (result.found) {
@@ -445,6 +479,8 @@ export const {
       }
     })
 
+    onCleanup(() => deltaBuffer.drain())
+
     const exit = useExit()
     const args = useArgs()