fix(ui): stabilize streaming message/tool rendering

Avoid remounting message blocks on part updates so tool call UI state persists. Render tool/message content from store and stabilize tool output scrolling during streaming.
2026-01-28 17:55:44 +00:00
parent d9bcc66930
commit a401eeec11
8 changed files with 289 additions and 100 deletions
--- a/packages/ui/src/components/message-block.tsx
+++ b/packages/ui/src/components/message-block.tsx
@@ -172,21 +172,212 @@ messageStoreBus.onInstanceDestroyed(clearInstanceCaches)
 interface ContentDisplayItem {
  type: "content"
  key: string
-  record: MessageRecord
-  parts: ClientPart[]
-  messageInfo?: MessageInfo
-  isQueued: boolean
-  showAgentMeta?: boolean
+  messageId: string
+  startPartId: string
 }

 interface ToolDisplayItem {
  type: "tool"
  key: string
-  toolPart: ToolCallPart
-  messageInfo?: MessageInfo
  messageId: string
-  messageVersion: number
-  partVersion: number
+  partId: string
+}
+
+interface MessageContentItemProps {
+  instanceId: string
+  sessionId: string
+  store: () => InstanceMessageStore
+  messageId: string
+  startPartId: string
+  messageIndex: number
+  lastAssistantIndex: () => number
+  onRevert?: (messageId: string) => void
+  onFork?: (messageId?: string) => void
+  onContentRendered?: () => void
+}
+
+function MessageContentItem(props: MessageContentItemProps) {
+  const record = createMemo(() => props.store().getMessage(props.messageId))
+  const messageInfo = createMemo(() => props.store().getMessageInfo(props.messageId))
+
+  const isQueued = createMemo(() => {
+    const current = record()
+    if (!current) return false
+    if (current.role !== "user") return false
+    const lastAssistant = props.lastAssistantIndex()
+    return lastAssistant === -1 || props.messageIndex > lastAssistant
+  })
+
+  const parts = createMemo<ClientPart[]>(() => {
+    const current = record()
+    if (!current) return []
+    const ids = current.partIds
+    const startIndex = ids.indexOf(props.startPartId)
+    if (startIndex === -1) return []
+
+    const resolved: ClientPart[] = []
+    for (let idx = startIndex; idx < ids.length; idx++) {
+      const partId = ids[idx]
+      const part = current.parts[partId]?.data
+      if (!part) continue
+      if (
+        part.type === "tool" ||
+        part.type === "reasoning" ||
+        part.type === "compaction" ||
+        part.type === "step-start" ||
+        part.type === "step-finish"
+      ) {
+        break
+      }
+      resolved.push(part)
+    }
+
+    return resolved
+  })
+
+  const showAgentMeta = createMemo(() => {
+    const current = record()
+    if (!current) return false
+    if (current.role !== "assistant") return false
+
+    const currentParts = parts()
+    if (!currentParts.some((part) => partHasRenderableText(part))) {
+      return false
+    }
+
+    const ids = current.partIds
+    const startIndex = ids.indexOf(props.startPartId)
+    if (startIndex === -1) return false
+
+    // Only show agent meta on the first content segment that contains renderable content.
+    for (let idx = 0; idx < startIndex; idx++) {
+      const partId = ids[idx]
+      const part = current.parts[partId]?.data
+      if (!part) continue
+      if (
+        part.type === "tool" ||
+        part.type === "reasoning" ||
+        part.type === "compaction" ||
+        part.type === "step-start" ||
+        part.type === "step-finish"
+      ) {
+        continue
+      }
+      if (partHasRenderableText(part)) {
+        return false
+      }
+    }
+
+    return true
+  })
+
+  return (
+    <Show when={record()}>
+      {(resolvedRecord) => (
+        <MessageItem
+          record={resolvedRecord()}
+          messageInfo={messageInfo()}
+          parts={parts()}
+          instanceId={props.instanceId}
+          sessionId={props.sessionId}
+          isQueued={isQueued()}
+          showAgentMeta={showAgentMeta()}
+          onRevert={props.onRevert}
+          onFork={props.onFork}
+          onContentRendered={props.onContentRendered}
+        />
+      )}
+    </Show>
+  )
+}
+
+interface ToolCallItemProps {
+  instanceId: string
+  sessionId: string
+  store: () => InstanceMessageStore
+  messageId: string
+  partId: string
+  onContentRendered?: () => void
+}
+
+function ToolCallItem(props: ToolCallItemProps) {
+  const { t } = useI18n()
+
+  const record = createMemo(() => props.store().getMessage(props.messageId))
+  const messageInfo = createMemo(() => props.store().getMessageInfo(props.messageId))
+  const partEntry = createMemo(() => record()?.parts?.[props.partId])
+
+  const toolPart = createMemo(() => {
+    const part = partEntry()?.data as ClientPart | undefined
+    if (!part || part.type !== "tool") return undefined
+    return part as ToolCallPart
+  })
+
+  const toolState = createMemo(() => toolPart()?.state as ToolState | undefined)
+  const toolName = createMemo(() => toolPart()?.tool || "")
+  const messageVersion = createMemo(() => record()?.revision ?? 0)
+  const partVersion = createMemo(() => partEntry()?.revision ?? 0)
+
+  const taskSessionId = createMemo(() => {
+    const state = toolState()
+    if (!state) return ""
+    if (!(isToolStateRunning(state) || isToolStateCompleted(state) || isToolStateError(state))) {
+      return ""
+    }
+    return extractTaskSessionId(state)
+  })
+
+  const taskLocation = createMemo(() => {
+    const id = taskSessionId()
+    if (!id) return null
+    return findTaskSessionLocation(id, props.instanceId)
+  })
+
+  const handleGoToTaskSession = (event: MouseEvent) => {
+    event.preventDefault()
+    event.stopPropagation()
+    const location = taskLocation()
+    if (!location) return
+    navigateToTaskSession(location)
+  }
+
+  return (
+    <Show when={toolPart()}>
+      {(resolvedToolPart) => (
+        <>
+          <div class="tool-call-header-label">
+            <div class="tool-call-header-meta">
+              <span class="tool-call-icon">{TOOL_ICON}</span>
+              <span>{t("messageBlock.tool.header")}</span>
+              <span class="tool-name">{toolName() || t("messageBlock.tool.unknown")}</span>
+            </div>
+            <Show when={taskSessionId()}>
+              <button
+                class="tool-call-header-button"
+                type="button"
+                disabled={!taskLocation()}
+                onClick={handleGoToTaskSession}
+                title={!taskLocation() ? t("messageBlock.tool.goToSession.unavailableTitle") : t("messageBlock.tool.goToSession.title")}
+              >
+                {t("messageBlock.tool.goToSession.label")}
+              </button>
+            </Show>
+          </div>
+
+          <ToolCall
+            toolCall={resolvedToolPart()}
+            toolCallId={props.partId}
+            messageId={props.messageId}
+            messageVersion={messageVersion()}
+            partVersion={partVersion()}
+            instanceId={props.instanceId}
+            sessionId={props.sessionId}
+            onContentRendered={props.onContentRendered}
+          />
+        </>
+      )}
+    </Show>
+  )
 }

 interface StepDisplayItem {
@@ -272,7 +463,6 @@ export default function MessageBlock(props: MessageBlockProps) {
    const items: MessageBlockItem[] = []
    const blockContentKeys: string[] = []
    const blockToolKeys: string[] = []
-    let segmentIndex = 0
    let pendingParts: ClientPart[] = []
    let agentMetaAttached = current.role !== "assistant"
    const defaultAccentColor = current.role === "user" ? USER_BORDER_COLOR : ASSISTANT_BORDER_COLOR
@@ -280,34 +470,28 @@ export default function MessageBlock(props: MessageBlockProps) {

    const flushContent = () => {
      if (pendingParts.length === 0) return
-      const segmentKey = `${current.id}:segment:${segmentIndex}`
-      segmentIndex += 1
-      const shouldShowAgentMeta =
-        current.role === "assistant" &&
-        !agentMetaAttached &&
-        pendingParts.some((part) => partHasRenderableText(part))
+      const startPartId = typeof (pendingParts[0] as any)?.id === "string" ? ((pendingParts[0] as any).id as string) : ""
+      if (!startPartId) {
+        pendingParts = []
+        return
+      }
+
+      if (!agentMetaAttached && pendingParts.some((part) => partHasRenderableText(part))) {
+        agentMetaAttached = true
+      }
+
+      const segmentKey = `${current.id}:content:${startPartId}`
      let cached = sessionCache.messageItems.get(segmentKey)
      if (!cached) {
        cached = {
          type: "content",
          key: segmentKey,
-          record: current,
-          parts: pendingParts.slice(),
-          messageInfo: info,
-          isQueued,
-          showAgentMeta: shouldShowAgentMeta,
+          messageId: current.id,
+          startPartId,
        }
        sessionCache.messageItems.set(segmentKey, cached)
-      } else {
-        cached.record = current
-        cached.parts = pendingParts.slice()
-        cached.messageInfo = info
-        cached.isQueued = isQueued
-        cached.showAgentMeta = shouldShowAgentMeta
-      }
-      if (shouldShowAgentMeta) {
-        agentMetaAttached = true
      }
+
      items.push(cached)
      blockContentKeys.push(segmentKey)
      lastAccentColor = defaultAccentColor
@@ -317,28 +501,26 @@ export default function MessageBlock(props: MessageBlockProps) {
    orderedParts.forEach((part, partIndex) => {
      if (part.type === "tool") {
        flushContent()
-        const partVersion = typeof (part as any).revision === "number" ? (part as any).revision : 0
-        const messageVersion = current.revision
-        const key = `${current.id}:${part.id ?? partIndex}`
+        const partId = part.id
+        if (!partId) {
+          // Tool parts are required to have ids; if one slips through, skip rendering
+          // to avoid unstable keys and accidental remount cascades.
+          return
+        }
+        const key = `${current.id}:${partId}`
        let toolItem = sessionCache.toolItems.get(key)
        if (!toolItem) {
          toolItem = {
            type: "tool",
            key,
-            toolPart: part as ToolCallPart,
-            messageInfo: info,
            messageId: current.id,
-            messageVersion,
-            partVersion,
+            partId,
          }
          sessionCache.toolItems.set(key, toolItem)
        } else {
          toolItem.key = key
-          toolItem.toolPart = part as ToolCallPart
-          toolItem.messageInfo = info
          toolItem.messageId = current.id
-          toolItem.messageVersion = messageVersion
-          toolItem.partVersion = partVersion
+          toolItem.partId = partId
        }
        items.push(toolItem)
        blockToolKeys.push(key)
@@ -427,21 +609,21 @@ export default function MessageBlock(props: MessageBlockProps) {
  })

  return (
-    <Show when={block()} keyed>
+    <Show when={block()}>
      {(resolvedBlock) => (
-        <div class="message-stream-block" data-message-id={resolvedBlock.record.id}>
-          <For each={resolvedBlock.items}>
+        <div class="message-stream-block" data-message-id={resolvedBlock().record.id}>
+          <For each={resolvedBlock().items}>
            {(item) => (
              <Switch>
                <Match when={item.type === "content"}>
-                  <MessageItem
-                    record={(item as ContentDisplayItem).record}
-                    messageInfo={(item as ContentDisplayItem).messageInfo}
-                    parts={(item as ContentDisplayItem).parts}
+                  <MessageContentItem
                    instanceId={props.instanceId}
                    sessionId={props.sessionId}
-                    isQueued={(item as ContentDisplayItem).isQueued}
-                    showAgentMeta={(item as ContentDisplayItem).showAgentMeta}
+                    store={props.store}
+                    messageId={(item as ContentDisplayItem).messageId}
+                    startPartId={(item as ContentDisplayItem).startPartId}
+                    messageIndex={props.messageIndex}
+                    lastAssistantIndex={props.lastAssistantIndex}
                    onRevert={props.onRevert}
                    onFork={props.onFork}
                    onContentRendered={props.onContentRendered}
@@ -450,46 +632,14 @@ export default function MessageBlock(props: MessageBlockProps) {
                <Match when={item.type === "tool"}>
                  {(() => {
                    const toolItem = item as ToolDisplayItem
-                    const toolState = toolItem.toolPart.state as ToolState | undefined
-                    const hasToolState =
-                      Boolean(toolState) && (isToolStateRunning(toolState) || isToolStateCompleted(toolState) || isToolStateError(toolState))
-                    const taskSessionId = hasToolState ? extractTaskSessionId(toolState) : ""
-                    const taskLocation = taskSessionId ? findTaskSessionLocation(taskSessionId, props.instanceId) : null
-                    const handleGoToTaskSession = (event: MouseEvent) => {
-                      event.preventDefault()
-                      event.stopPropagation()
-                      if (!taskLocation) return
-                      navigateToTaskSession(taskLocation)
-                    }
-
                    return (
                      <div class="tool-call-message" data-key={toolItem.key}>
-                        <div class="tool-call-header-label">
-                          <div class="tool-call-header-meta">
-                            <span class="tool-call-icon">{TOOL_ICON}</span>
-                            <span>{t("messageBlock.tool.header")}</span>
-                            <span class="tool-name">{toolItem.toolPart.tool || t("messageBlock.tool.unknown")}</span>
-                          </div>
-                          <Show when={taskSessionId}>
-                            <button
-                              class="tool-call-header-button"
-                              type="button"
-                              disabled={!taskLocation}
-                              onClick={handleGoToTaskSession}
-                              title={!taskLocation ? t("messageBlock.tool.goToSession.unavailableTitle") : t("messageBlock.tool.goToSession.title")}
-                            >
-                              {t("messageBlock.tool.goToSession.label")}
-                            </button>
-                          </Show>
-                        </div>
-                        <ToolCall
-                          toolCall={toolItem.toolPart}
-                          toolCallId={toolItem.toolPart.id}
-                          messageId={toolItem.messageId}
-                          messageVersion={toolItem.messageVersion}
-                          partVersion={toolItem.partVersion}
+                        <ToolCallItem
                          instanceId={props.instanceId}
                          sessionId={props.sessionId}
+                          store={props.store}
+                          messageId={toolItem.messageId}
+                          partId={toolItem.partId}
                          onContentRendered={props.onContentRendered}
                        />
                      </div>
--- a/packages/ui/src/components/message-item.tsx
+++ b/packages/ui/src/components/message-item.tsx
@@ -137,8 +137,17 @@ export default function MessageItem(props: MessageItemProps) {
  }

  const isGenerating = () => {
+    if (hasContent()) {
+      return false
+    }
+
+    // Prefer the local record status for streaming placeholders.
+    if (!isUser() && props.record.status === "streaming") {
+      return true
+    }
+
    const info = props.messageInfo
-    return !hasContent() && info && info.role === "assistant" && info.time.completed !== undefined && info.time.completed === 0
+    return Boolean(info && info.role === "assistant" && info.time.completed !== undefined && info.time.completed === 0)
  }

  const handleRevert = () => {
@@ -163,7 +172,7 @@ export default function MessageItem(props: MessageItemProps) {
    setTimeout(() => setCopied(false), 2000)
  }

-  if (!isUser() && !hasContent()) {
+  if (!isUser() && !hasContent() && !isGenerating()) {
    return null
  }

--- a/packages/ui/src/components/message-part.tsx
+++ b/packages/ui/src/components/message-part.tsx
@@ -25,6 +25,13 @@ interface MessagePartProps {
  const isAssistantMessage = () => props.messageType === "assistant"
  const textContainerClass = () => (isAssistantMessage() ? "message-text message-text-assistant" : "message-text")

+  const shouldHideTextPart = () => {
+    const part = props.part
+    if (!part || part.type !== "text") return false
+    // Keep optimistic user prompts visible; hide synthetic assistant text.
+    return Boolean((part as any).synthetic) && props.messageType !== "user"
+  }
+
  const plainTextContent = () => {
    const part = props.part

@@ -94,7 +101,7 @@ interface MessagePartProps {
  return (
    <Switch>
      <Match when={partType() === "text"}>
-        <Show when={!(props.part.type === "text" && props.part.synthetic) && partHasRenderableText(props.part)}>
+        <Show when={!shouldHideTextPart() && partHasRenderableText(props.part)}>
          <div class={textContainerClass()}>
            <Show
               when={isAssistantMessage()}
--- a/packages/ui/src/components/tool-call.tsx
+++ b/packages/ui/src/components/tool-call.tsx
@@ -235,12 +235,16 @@ export default function ToolCall(props: ToolCallProps) {
    requestAnimationFrame(() => {
      restoreScrollPosition(autoScroll())
      if (!expanded()) return
-      scheduleAnchorScroll()
+      scheduleAnchorScroll(true)
    })
  }

  const initializeScrollContainer = (element: HTMLDivElement | null | undefined) => {
-    scrollContainerRef = element || undefined
+    const next = element || undefined
+    if (next === scrollContainerRef) {
+      return
+    }
+    scrollContainerRef = next
    setScrollContainer(scrollContainerRef)
    if (scrollContainerRef) {
      restoreScrollPosition(autoScroll())
@@ -593,7 +597,7 @@ export default function ToolCall(props: ToolCallProps) {
      return
    }
    previousPartVersion = version
-    scheduleAnchorScroll()
+    scheduleAnchorScroll(true)
  })

  createEffect(() => {
--- a/packages/ui/src/components/tool-call/ansi-render.tsx
+++ b/packages/ui/src/components/tool-call/ansi-render.tsx
@@ -87,7 +87,7 @@ export function createAnsiContentRenderer(params: {
    }

    return (
-      <div class={messageClass} ref={(element) => params.scrollHelpers.registerContainer(element)} onScroll={params.scrollHelpers.handleScroll}>
+      <div class={messageClass} ref={params.scrollHelpers.registerContainer} onScroll={params.scrollHelpers.handleScroll}>
        <pre class="tool-call-content tool-call-ansi" innerHTML={nextCache.html} />
        {params.scrollHelpers.renderSentinel()}
      </div>
--- a/packages/ui/src/components/tool-call/diff-render.tsx
+++ b/packages/ui/src/components/tool-call/diff-render.tsx
@@ -26,6 +26,14 @@ export function createDiffContentRenderer(params: {
  handleScrollRendered: () => void
  onContentRendered?: () => void
 }) {
+  const registerTracked = (element: HTMLDivElement | null) => {
+    params.scrollHelpers.registerContainer(element)
+  }
+
+  const registerUntracked = (element: HTMLDivElement | null) => {
+    params.scrollHelpers.registerContainer(element, { disableTracking: true })
+  }
+
  function renderDiffContent(payload: DiffPayload, options?: DiffRenderOptions): JSXElement | null {
    const relativePath = payload.filePath ? getRelativePath(payload.filePath) : ""
    const toolbarLabel = options?.label || (relativePath
@@ -35,6 +43,8 @@ export function createDiffContentRenderer(params: {
    const cacheHandle = selectedVariant === "permission-diff" ? params.permissionDiffCache : params.diffCache
    const diffMode = () => (params.preferences().diffViewMode || "split") as DiffViewMode
    const themeKey = params.isDark() ? "dark" : "light"
+    const disableScrollTracking = Boolean(options?.disableScrollTracking)
+    const registerRef = disableScrollTracking ? registerUntracked : registerTracked

    const baseEntryParams = cacheHandle.params() as any
    const cacheEntryParams = (() => {
@@ -58,7 +68,7 @@ export function createDiffContentRenderer(params: {
    }

    const handleDiffRendered = () => {
-      if (!options?.disableScrollTracking) {
+      if (!disableScrollTracking) {
        params.handleScrollRendered()
      }
      params.onContentRendered?.()
@@ -67,8 +77,8 @@ export function createDiffContentRenderer(params: {
    return (
      <div
        class="message-text tool-call-markdown tool-call-markdown-large tool-call-diff-shell"
-        ref={(element) => params.scrollHelpers.registerContainer(element, { disableTracking: options?.disableScrollTracking })}
-        onScroll={options?.disableScrollTracking ? undefined : params.scrollHelpers.handleScroll}
+        ref={registerRef}
+        onScroll={disableScrollTracking ? undefined : params.scrollHelpers.handleScroll}
      >
        <div class="tool-call-diff-toolbar" role="group" aria-label={params.t("toolCall.diff.viewMode.ariaLabel")}>
          <span class="tool-call-diff-toolbar-label">{toolbarLabel}</span>
@@ -100,7 +110,7 @@ export function createDiffContentRenderer(params: {
          cacheEntryParams={cacheEntryParams as any}
          onRendered={handleDiffRendered}
        />
-        {params.scrollHelpers.renderSentinel({ disableTracking: options?.disableScrollTracking })}
+        {params.scrollHelpers.renderSentinel({ disableTracking: disableScrollTracking })}
      </div>
    )
  }
--- a/packages/ui/src/components/tool-call/markdown-render.tsx
+++ b/packages/ui/src/components/tool-call/markdown-render.tsx
@@ -15,6 +15,14 @@ export function createMarkdownContentRenderer(params: {
  handleScrollRendered: () => void
  onContentRendered?: () => void
 }) {
+  const registerTracked = (element: HTMLDivElement | null) => {
+    params.scrollHelpers.registerContainer(element)
+  }
+
+  const registerUntracked = (element: HTMLDivElement | null) => {
+    params.scrollHelpers.registerContainer(element, { disableTracking: true })
+  }
+
  function renderMarkdownContent(options: MarkdownRenderOptions): JSXElement | null {
    if (!options.content) {
      return null
@@ -24,6 +32,7 @@ export function createMarkdownContentRenderer(params: {
    const disableHighlight = options.disableHighlight || false
    const messageClass = `message-text tool-call-markdown${size === "large" ? " tool-call-markdown-large" : ""}`
    const disableScrollTracking = options.disableScrollTracking || false
+    const registerRef = disableScrollTracking ? registerUntracked : registerTracked

    const state = params.toolState()
    const shouldDeferMarkdown = Boolean(state && (state.status === "running" || state.status === "pending") && disableHighlight)
@@ -31,7 +40,7 @@ export function createMarkdownContentRenderer(params: {
      return (
        <div
          class={messageClass}
-          ref={(element) => params.scrollHelpers.registerContainer(element, { disableTracking: disableScrollTracking })}
+          ref={registerRef}
          onScroll={disableScrollTracking ? undefined : params.scrollHelpers.handleScroll}
        >
          <pre class="whitespace-pre-wrap break-words text-sm font-mono">{options.content}</pre>
@@ -56,7 +65,7 @@ export function createMarkdownContentRenderer(params: {
    return (
      <div
        class={messageClass}
-        ref={(element) => params.scrollHelpers.registerContainer(element, { disableTracking: disableScrollTracking })}
+        ref={registerRef}
        onScroll={disableScrollTracking ? undefined : params.scrollHelpers.handleScroll}
      >
        <Markdown
--- a/packages/ui/src/components/tool-call/renderers/task.tsx
+++ b/packages/ui/src/components/tool-call/renderers/task.tsx
@@ -176,7 +176,7 @@ export const taskRenderer: ToolRenderer = {
            <div class="tool-call-task-section-body">
              <div
                class="message-text tool-call-markdown tool-call-task-container"
-                ref={(element) => scrollHelpers?.registerContainer(element)}
+                ref={scrollHelpers?.registerContainer}
                onScroll={
                  scrollHelpers ? (event) => scrollHelpers.handleScroll(event as Event & { currentTarget: HTMLDivElement }) : undefined
                }