From 6597783e850e10d75e55e1c0c5e2cec0ebe94fd3 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Tue, 28 Oct 2025 13:48:56 +0000 Subject: [PATCH] Stabilize message stream rendering caches --- src/components/message-item.tsx | 7 +- src/components/message-stream.tsx | 165 ++++++++++++++++++++++-------- src/stores/sessions.ts | 70 +++++++++++-- src/types/message.ts | 4 + 4 files changed, 190 insertions(+), 56 deletions(-) diff --git a/src/components/message-item.tsx b/src/components/message-item.tsx index 7acd15e2..6782eaea 100644 --- a/src/components/message-item.tsx +++ b/src/components/message-item.tsx @@ -6,6 +6,7 @@ interface MessageItemProps { message: Message messageInfo?: any isQueued?: boolean + parts?: any[] onRevert?: (messageId: string) => void } @@ -16,6 +17,8 @@ export default function MessageItem(props: MessageItemProps) { return date.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" }) } + const messageParts = () => props.parts ?? props.message.parts + const errorMessage = () => { if (!props.messageInfo?.error) return null @@ -36,7 +39,7 @@ export default function MessageItem(props: MessageItemProps) { } const hasContent = () => { - return props.message.parts.length > 0 || errorMessage() !== null + return messageParts().length > 0 || errorMessage() !== null } const isGenerating = () => { @@ -81,7 +84,7 @@ export default function MessageItem(props: MessageItemProps) { - {(part) => } + {(part) => } diff --git a/src/components/message-stream.tsx b/src/components/message-stream.tsx index 3980df60..d0db98d6 100644 --- a/src/components/message-stream.tsx +++ b/src/components/message-stream.tsx @@ -1,11 +1,11 @@ import { For, Show, createSignal, createEffect, createMemo } from "solid-js" -import type { Message } from "../types/message" +import type { Message, MessageDisplayParts } from "../types/message" import MessageItem from "./message-item" import ToolCall from "./tool-call" import { sseManager } from "../lib/sse-manager" import Kbd from "./kbd" import { preferences } from "../stores/preferences" -import { providers, getSessionInfo } from "../stores/sessions" +import { providers, getSessionInfo, computeDisplayParts } from "../stores/sessions" // Calculate session tokens and cost from messagesInfo (matches TUI logic) function calculateSessionInfo(messagesInfo?: Map, instanceId?: string) { @@ -120,17 +120,46 @@ interface MessageStreamProps { onRevert?: (messageId: string) => void } -interface DisplayItem { - type: "message" | "tool" - data: any +interface MessageDisplayItem { + type: "message" + message: Message + combinedParts: any[] + isQueued: boolean messageInfo?: any } +interface ToolDisplayItem { + type: "tool" + key: string + toolPart: any + messageInfo?: any +} + +type DisplayItem = MessageDisplayItem | ToolDisplayItem + +interface MessageCacheEntry { + version: number + showThinking: boolean + isQueued: boolean + messageInfo?: any + displayParts: MessageDisplayParts + item: MessageDisplayItem +} + +interface ToolCacheEntry { + toolPart: any + messageInfo?: any + item: ToolDisplayItem +} + export default function MessageStream(props: MessageStreamProps) { let containerRef: HTMLDivElement | undefined const [autoScroll, setAutoScroll] = createSignal(true) const [showScrollButton, setShowScrollButton] = createSignal(false) + let messageItemCache = new Map() + let toolItemCache = new Map() + const connectionStatus = () => sseManager.getStatus(props.instanceId) const sessionInfo = createMemo(() => { @@ -180,9 +209,11 @@ export default function MessageStream(props: MessageStreamProps) { const displayItems = createMemo(() => { // Ensure memo reacts to preference changes - preferences().showThinkingBlocks + const showThinking = preferences().showThinkingBlocks const items: DisplayItem[] = [] + const newMessageCache = new Map() + const newToolCache = new Map() let lastAssistantIndex = -1 for (let i = props.messages.length - 1; i >= 0; i--) { @@ -201,35 +232,82 @@ export default function MessageStream(props: MessageStreamProps) { break } - // Use precomputed displayParts, fallback to empty arrays if not available - const displayParts = message.displayParts || { text: [], tool: [], reasoning: [] } - const textParts = displayParts.text - const toolParts = displayParts.tool - const reasoningParts = displayParts.reasoning + const baseDisplayParts = message.displayParts + const displayParts: MessageDisplayParts = + baseDisplayParts && baseDisplayParts.showThinking === showThinking + ? baseDisplayParts + : computeDisplayParts(message, showThinking) + const combinedParts = displayParts.combined + const version = message.version ?? 0 const isQueued = message.type === "user" && (lastAssistantIndex === -1 || index > lastAssistantIndex) - if (textParts.length > 0 || reasoningParts.length > 0 || messageInfo?.error) { - items.push({ + const cacheEntry = messageItemCache.get(message.id) + if ( + cacheEntry && + cacheEntry.version === version && + cacheEntry.showThinking === showThinking && + cacheEntry.isQueued === isQueued && + cacheEntry.messageInfo === messageInfo + ) { + cacheEntry.displayParts = displayParts + cacheEntry.version = version + cacheEntry.showThinking = showThinking + cacheEntry.isQueued = isQueued + cacheEntry.messageInfo = messageInfo + cacheEntry.item.message = message + cacheEntry.item.combinedParts = combinedParts + cacheEntry.item.isQueued = isQueued + cacheEntry.item.messageInfo = messageInfo + newMessageCache.set(message.id, cacheEntry) + items.push(cacheEntry.item) + } else { + const messageItem: MessageDisplayItem = { type: "message", - data: { - ...message, - parts: [...textParts, ...reasoningParts], - isQueued, - }, + message, + combinedParts, + isQueued, messageInfo, + } + newMessageCache.set(message.id, { + version, + showThinking, + isQueued, + messageInfo, + displayParts, + item: messageItem, }) + items.push(messageItem) } - for (const toolPart of toolParts) { - items.push({ - type: "tool", - data: toolPart, - messageInfo, - }) + for (let toolIndex = 0; toolIndex < displayParts.tool.length; toolIndex++) { + const toolPart = displayParts.tool[toolIndex] + const toolKey = typeof toolPart?.id === "string" ? toolPart.id : `${message.id}-tool-${toolIndex}` + + const toolEntry = toolItemCache.get(toolKey) + if (toolEntry && toolEntry.toolPart === toolPart && toolEntry.messageInfo === messageInfo) { + toolEntry.item.toolPart = toolPart + toolEntry.item.messageInfo = messageInfo + toolEntry.toolPart = toolPart + toolEntry.messageInfo = messageInfo + newToolCache.set(toolKey, toolEntry) + items.push(toolEntry.item) + } else { + const toolItem: ToolDisplayItem = { + type: "tool", + key: toolKey, + toolPart, + messageInfo, + } + newToolCache.set(toolKey, { toolPart, messageInfo, item: toolItem }) + items.push(toolItem) + } } } + messageItemCache = newMessageCache + toolItemCache = newToolCache + return items }) @@ -301,29 +379,30 @@ export default function MessageStream(props: MessageStreamProps) { - {(item, index) => { - const key = item.type === "message" ? `msg-${item.data.id}` : `tool-${item.data.id}` - return ( - -
- 🔧 - Tool Call - {item.data?.tool || "unknown"} -
- - - } - > + {(item) => { + if (item.type === "message") { + return ( -
+ ) + } + + const toolPart = item.toolPart + + return ( +
+
+ 🔧 + Tool Call + {toolPart?.tool || "unknown"} +
+ +
) }}
diff --git a/src/stores/sessions.ts b/src/stores/sessions.ts index 8ddababe..1fddfae0 100644 --- a/src/stores/sessions.ts +++ b/src/stores/sessions.ts @@ -82,7 +82,7 @@ function removeSessionIndexes(instanceId: string) { sessionIndexes.delete(instanceId) } -function computeDisplayParts(message: Message, showThinking: boolean): MessageDisplayParts { +export function computeDisplayParts(message: Message, showThinking: boolean): MessageDisplayParts { const text: any[] = [] const tool: any[] = [] const reasoning: any[] = [] @@ -97,7 +97,10 @@ function computeDisplayParts(message: Message, showThinking: boolean): MessageDi } } - return { text, tool, reasoning } + const combined = reasoning.length > 0 ? [...text, ...reasoning] : [...text] + const version = typeof message.version === "number" ? message.version : 0 + + return { text, tool, reasoning, combined, showThinking, version } } function withSession(instanceId: string, sessionId: string, updater: (session: Session) => void) { @@ -710,6 +713,7 @@ async function loadMessages(instanceId: string, sessionId: string, force = false parts: apiMessage.parts || [], timestamp: info.time?.created || Date.now(), status: "complete" as const, + version: 0, } message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) @@ -824,6 +828,7 @@ function handleMessageUpdate(instanceId: string, event: any): void { parts: [part], timestamp: Date.now(), status: "streaming" as const, + version: 0, } newMessage.displayParts = computeDisplayParts(newMessage, preferences().showThinkingBlocks) @@ -841,6 +846,9 @@ function handleMessageUpdate(instanceId: string, event: any): void { } else { // Update existing message const message = session.messages[messageIndex] + if (typeof message.version !== "number") { + message.version = 0 + } // Strip synthetic parts when real data arrives let filteredSynthetics = false @@ -864,14 +872,32 @@ function handleMessageUpdate(instanceId: string, event: any): void { index.partIndex.set(message.id, partMap) } + let shouldIncrementVersion = filteredSynthetics || replacedTemp const partIndex = partMap.get(part.id) + if (partIndex === undefined) { baseParts.push(part) if (part.id && typeof part.id === "string") { partMap.set(part.id, baseParts.length - 1) } + shouldIncrementVersion = true } else { + const previousPart = baseParts[partIndex] + const textUnchanged = + !filteredSynthetics && + !replacedTemp && + part.type === "text" && + previousPart?.type === "text" && + previousPart.text === part.text + + if (textUnchanged) { + return + } + baseParts[partIndex] = part + if (part.type !== "text" || !previousPart || previousPart.text !== part.text) { + shouldIncrementVersion = true + } } const oldId = message.id @@ -879,7 +905,16 @@ function handleMessageUpdate(instanceId: string, event: any): void { message.status = message.status === "sending" ? "streaming" : message.status message.parts = baseParts - message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) + if (shouldIncrementVersion) { + message.version += 1 + message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) + } else if ( + !message.displayParts || + message.displayParts.showThinking !== preferences().showThinkingBlocks || + message.displayParts.version !== message.version + ) { + message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) + } // Update message index if ID changed if (oldId !== message.id) { @@ -947,11 +982,17 @@ function handleMessageUpdate(instanceId: string, event: any): void { if (tempMessageIndex > -1) { // Replace queued message const message = session.messages[tempMessageIndex] + if (typeof message.version !== "number") { + message.version = 0 + } + const oldId = message.id message.id = info.id message.type = (info.role === "user" ? "user" : "assistant") as "user" | "assistant" message.timestamp = info.time?.created || Date.now() message.status = "complete" as const + message.version += 1 + message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) if (oldId !== message.id) { index.messageIndex.delete(oldId) @@ -971,6 +1012,7 @@ function handleMessageUpdate(instanceId: string, event: any): void { parts: [], timestamp: info.time?.created || Date.now(), status: "complete" as const, + version: 0, } newMessage.displayParts = computeDisplayParts(newMessage, preferences().showThinkingBlocks) @@ -989,16 +1031,21 @@ function handleMessageUpdate(instanceId: string, event: any): void { } else { // Update existing message status const message = session.messages[messageIndex] + if (typeof message.version !== "number") { + message.version = 0 + } message.status = "complete" as const + message.version += 1 + message.displayParts = computeDisplayParts(message, preferences().showThinkingBlocks) + + session.messagesInfo.set(info.id, info) + + withSession(instanceId, info.sessionID, (session) => { + // Session already mutated in place + }) + + updateSessionInfo(instanceId, info.sessionID) } - - session.messagesInfo.set(info.id, info) - - withSession(instanceId, info.sessionID, (session) => { - // Session already mutated in place - }) - - updateSessionInfo(instanceId, info.sessionID) } } @@ -1110,6 +1157,7 @@ async function sendMessage( parts: optimisticParts, timestamp: Date.now(), status: "sending", + version: 0, } optimisticMessage.displayParts = computeDisplayParts(optimisticMessage, preferences().showThinkingBlocks) diff --git a/src/types/message.ts b/src/types/message.ts index 01dbe187..c2a54ec8 100644 --- a/src/types/message.ts +++ b/src/types/message.ts @@ -2,6 +2,9 @@ export interface MessageDisplayParts { text: any[] tool: any[] reasoning: any[] + combined: any[] + showThinking: boolean + version: number } export interface Message { @@ -11,5 +14,6 @@ export interface Message { parts: any[] timestamp: number status: "sending" | "sent" | "streaming" | "complete" | "error" + version: number displayParts?: MessageDisplayParts }