From d0934545b1ad1b6b093a93eabadc184adee1aca9 Mon Sep 17 00:00:00 2001 From: kronos Date: Mon, 31 Aug 2026 19:16:41 -0400 Subject: [PATCH] fix(tui): coalesce message.part.delta store writes (V1 path) --- packages/tui/src/context/delta-buffer.test.ts | 111 ++++++++++++++++++ packages/tui/src/context/delta-buffer.ts | 53 +++++++++ packages/tui/src/context/sync.tsx | 58 +++++++-- 3 files changed, 211 insertions(+), 11 deletions(-) create mode 100644 packages/tui/src/context/delta-buffer.test.ts create mode 100644 packages/tui/src/context/delta-buffer.ts diff --git a/packages/tui/src/context/delta-buffer.test.ts b/packages/tui/src/context/delta-buffer.test.ts new file mode 100644 index 000000000000..51233cb5407f --- /dev/null +++ b/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([]) + }) +}) diff --git a/packages/tui/src/context/delta-buffer.ts b/packages/tui/src/context/delta-buffer.ts new file mode 100644 index 000000000000..d912d85b8af9 --- /dev/null +++ b/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() + 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) + } + }, + } +} diff --git a/packages/tui/src/context/sync.tsx b/packages/tui/src/context/sync.tsx index 71e050d11e68..3c3f4b2c689c 100644 --- a/packages/tui/src/context/sync.tsx +++ b/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()