From a401eeec11b89be395499dd2b5431faac23a8259 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Wed, 28 Jan 2026 17:55:44 +0000 Subject: [PATCH] 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. --- packages/ui/src/components/message-block.tsx | 322 +++++++++++++----- packages/ui/src/components/message-item.tsx | 13 +- packages/ui/src/components/message-part.tsx | 9 +- packages/ui/src/components/tool-call.tsx | 10 +- .../src/components/tool-call/ansi-render.tsx | 2 +- .../src/components/tool-call/diff-render.tsx | 18 +- .../components/tool-call/markdown-render.tsx | 13 +- .../components/tool-call/renderers/task.tsx | 2 +- 8 files changed, 289 insertions(+), 100 deletions(-) diff --git a/packages/ui/src/components/message-block.tsx b/packages/ui/src/components/message-block.tsx index b7b47820..a9a01eea 100644 --- 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(() => { + 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 ( + + {(resolvedRecord) => ( + + )} + + ) +} + +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 ( + + {(resolvedToolPart) => ( + <> +
+
+ {TOOL_ICON} + {t("messageBlock.tool.header")} + {toolName() || t("messageBlock.tool.unknown")} +
+ + + +
+ + + + )} + + ) } 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 ( - + {(resolvedBlock) => ( -
- +
+ {(item) => ( - {(() => { 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 (
-
-
- {TOOL_ICON} - {t("messageBlock.tool.header")} - {toolItem.toolPart.tool || t("messageBlock.tool.unknown")} -
- - - -
-
diff --git a/packages/ui/src/components/message-item.tsx b/packages/ui/src/components/message-item.tsx index 7f3cb0a7..e066e34e 100644 --- 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 } diff --git a/packages/ui/src/components/message-part.tsx b/packages/ui/src/components/message-part.tsx index 421985bf..50e7759f 100644 --- 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 ( - +
{ 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(() => { diff --git a/packages/ui/src/components/tool-call/ansi-render.tsx b/packages/ui/src/components/tool-call/ansi-render.tsx index 2d7a49c1..8a5ed099 100644 --- 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 ( -
params.scrollHelpers.registerContainer(element)} onScroll={params.scrollHelpers.handleScroll}> + 
         {params.scrollHelpers.renderSentinel()}
diff --git a/packages/ui/src/components/tool-call/diff-render.tsx b/packages/ui/src/components/tool-call/diff-render.tsx index 99113ef7..4f215c39 100644 --- 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 (
params.scrollHelpers.registerContainer(element, { disableTracking: options?.disableScrollTracking })} - onScroll={options?.disableScrollTracking ? undefined : params.scrollHelpers.handleScroll} + ref={registerRef} + onScroll={disableScrollTracking ? undefined : params.scrollHelpers.handleScroll} >
{toolbarLabel} @@ -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 })}
) } diff --git a/packages/ui/src/components/tool-call/markdown-render.tsx b/packages/ui/src/components/tool-call/markdown-render.tsx index 36969901..c94859db 100644 --- 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 (
params.scrollHelpers.registerContainer(element, { disableTracking: disableScrollTracking })} + ref={registerRef} onScroll={disableScrollTracking ? undefined : params.scrollHelpers.handleScroll} > 
{options.content}
@@ -56,7 +65,7 @@ export function createMarkdownContentRenderer(params: { return (
params.scrollHelpers.registerContainer(element, { disableTracking: disableScrollTracking })} + ref={registerRef} onScroll={disableScrollTracking ? undefined : params.scrollHelpers.handleScroll} >
scrollHelpers?.registerContainer(element)} + ref={scrollHelpers?.registerContainer} onScroll={ scrollHelpers ? (event) => scrollHelpers.handleScroll(event as Event & { currentTarget: HTMLDivElement }) : undefined }