From 94233261938ca4c28649c6199814bd46ea647c51 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Sun, 30 Nov 2025 20:56:17 +0000
Subject: [PATCH] Make MessageBlock rerender via keyed Show

---
 .../ui/src/components/message-stream-v2.tsx   | 197 +++++++++---------
 1 file changed, 99 insertions(+), 98 deletions(-)

diff --git a/packages/ui/src/components/message-stream-v2.tsx b/packages/ui/src/components/message-stream-v2.tsx
index 42e43b4a..4f56cc96 100644
--- a/packages/ui/src/components/message-stream-v2.tsx
+++ b/packages/ui/src/components/message-stream-v2.tsx
@@ -883,107 +883,108 @@ function MessageBlock(props: MessageBlockProps) {
     return resultBlock
   })
 
-  const resolvedBlock = block()
-  if (!resolvedBlock) return null
-
   return (
-    <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}
-                combinedParts={(item as ContentDisplayItem).parts}
-                orderedParts={(item as ContentDisplayItem).parts}
-                instanceId={props.instanceId}
-                sessionId={props.sessionId}
-                isQueued={(item as ContentDisplayItem).isQueued}
-                showAgentMeta={(item as ContentDisplayItem).showAgentMeta}
-                onRevert={props.onRevert}
-                onFork={props.onFork}
-              />
-            </Match>
-            <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) : null
-                const handleGoToTaskSession = (event: MouseEvent) => {
-                  event.preventDefault()
-                  event.stopPropagation()
-                  if (!taskLocation) return
-                  navigateToTaskSession(taskLocation)
-                }
+    <Show when={block()} keyed>
+      {(resolvedBlock) => (
+        <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}
+                    combinedParts={(item as ContentDisplayItem).parts}
+                    orderedParts={(item as ContentDisplayItem).parts}
+                    instanceId={props.instanceId}
+                    sessionId={props.sessionId}
+                    isQueued={(item as ContentDisplayItem).isQueued}
+                    showAgentMeta={(item as ContentDisplayItem).showAgentMeta}
+                    onRevert={props.onRevert}
+                    onFork={props.onFork}
+                  />
+                </Match>
+                <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) : 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>Tool Call</span>
-                        <span class="tool-name">{toolItem.toolPart.tool || "unknown"}</span>
+                    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>Tool Call</span>
+                            <span class="tool-name">{toolItem.toolPart.tool || "unknown"}</span>
+                          </div>
+                          <Show when={taskSessionId}>
+                            <button
+                              class="tool-call-header-button"
+                              type="button"
+                              disabled={!taskLocation}
+                              onClick={handleGoToTaskSession}
+                              title={!taskLocation ? "Session not available yet" : "Go to session"}
+                            >
+                              Go to Session
+                            </button>
+                          </Show>
+                        </div>
+                        <ToolCall
+                          toolCall={toolItem.toolPart}
+                          toolCallId={toolItem.key}
+                          messageId={toolItem.messageId}
+                          messageVersion={toolItem.messageVersion}
+                          partVersion={toolItem.partVersion}
+                          instanceId={props.instanceId}
+                          sessionId={props.sessionId}
+                        />
                       </div>
-                      <Show when={taskSessionId}>
-                        <button
-                          class="tool-call-header-button"
-                          type="button"
-                          disabled={!taskLocation}
-                          onClick={handleGoToTaskSession}
-                          title={!taskLocation ? "Session not available yet" : "Go to session"}
-                        >
-                          Go to Session
-                        </button>
-                      </Show>
-                    </div>
-                    <ToolCall
-                      toolCall={toolItem.toolPart}
-                      toolCallId={toolItem.key}
-                      messageId={toolItem.messageId}
-                      messageVersion={toolItem.messageVersion}
-                      partVersion={toolItem.partVersion}
-                      instanceId={props.instanceId}
-                      sessionId={props.sessionId}
-                    />
-                  </div>
-                )
-              })()}
-            </Match>
-            <Match when={item.type === "step-start"}>
-              <StepCard
-                kind="start"
-                part={(item as StepDisplayItem).part}
-                messageInfo={(item as StepDisplayItem).messageInfo}
-                showAgentMeta
-              />
-            </Match>
-            <Match when={item.type === "step-finish"}>
-              <StepCard
-                kind="finish"
-                part={(item as StepDisplayItem).part}
-                messageInfo={(item as StepDisplayItem).messageInfo}
-                showUsage={props.showUsageMetrics()}
-                borderColor={(item as StepDisplayItem).accentColor}
-              />
-            </Match>
-            <Match when={item.type === "reasoning"}>
-              <ReasoningCard
-                part={(item as ReasoningDisplayItem).part}
-                messageInfo={(item as ReasoningDisplayItem).messageInfo}
-                instanceId={props.instanceId}
-                sessionId={props.sessionId}
-                showAgentMeta={(item as ReasoningDisplayItem).showAgentMeta}
-                defaultExpanded={(item as ReasoningDisplayItem).defaultExpanded}
-              />
-            </Match>
-          </Switch>
-        )}
-      </For>
-    </div>
+                    )
+                  })()}
+                </Match>
+                <Match when={item.type === "step-start"}>
+                  <StepCard
+                    kind="start"
+                    part={(item as StepDisplayItem).part}
+                    messageInfo={(item as StepDisplayItem).messageInfo}
+                    showAgentMeta
+                  />
+                </Match>
+                <Match when={item.type === "step-finish"}>
+                  <StepCard
+                    kind="finish"
+                    part={(item as StepDisplayItem).part}
+                    messageInfo={(item as StepDisplayItem).messageInfo}
+                    showUsage={props.showUsageMetrics()}
+                    borderColor={(item as StepDisplayItem).accentColor}
+                  />
+                </Match>
+                <Match when={item.type === "reasoning"}>
+                  <ReasoningCard
+                    part={(item as ReasoningDisplayItem).part}
+                    messageInfo={(item as ReasoningDisplayItem).messageInfo}
+                    instanceId={props.instanceId}
+                    sessionId={props.sessionId}
+                    showAgentMeta={(item as ReasoningDisplayItem).showAgentMeta}
+                    defaultExpanded={(item as ReasoningDisplayItem).defaultExpanded}
+                  />
+                </Match>
+              </Switch>
+            )}
+          </For>
+        </div>
+      )}
+    </Show>
   )
 }