diff --git a/packages/ui/src/components/message-stream-v2.tsx b/packages/ui/src/components/message-stream-v2.tsx
index e1673d41..26e178e2 100644
--- a/packages/ui/src/components/message-stream-v2.tsx
+++ b/packages/ui/src/components/message-stream-v2.tsx
@@ -225,6 +225,20 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
   const showUsagePreference = () => preferences().showUsageMetrics ?? true
   const store = createMemo(() => messageStoreBus.getOrCreate(props.instanceId))
   const messageIds = createMemo(() => store().getSessionMessageIds(props.sessionId))
+  const visibleMessageIds = createMemo(() => {
+    const ids = messageIds()
+    const revert = store().getSessionRevert(props.sessionId)
+    if (!revert?.messageID) {
+      return ids
+    }
+    const stopIndex = ids.indexOf(revert.messageID)
+    return stopIndex === -1 ? ids : ids.slice(0, stopIndex)
+  })
+  const messageRecords = createMemo(() =>
+    visibleMessageIds()
+      .map((id) => store().getMessage(id))
+      .filter((record): record is MessageRecord => Boolean(record)),
+  )
 
   const sessionRevision = createMemo(() => store().getSessionRevision(props.sessionId))
   const usageSnapshot = createMemo(() => store().getSessionUsage(props.sessionId))
@@ -696,9 +710,7 @@ function MessageBlock(props: MessageBlockProps) {
 
   const block = createMemo<MessageDisplayBlock | null>(() => {
     const current = record()
-    if (!current) {
-      return null
-    }
+    if (!current) return null
 
     const index = props.messageIndexMap.get(current.id) ?? 0
     const isQueued = current.role === "user" && (props.lastAssistantIndex === -1 || index > props.lastAssistantIndex)
@@ -870,108 +882,107 @@ function MessageBlock(props: MessageBlockProps) {
     return resultBlock
   })
 
-  return (
-    <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)
-                    }
+  const resolvedBlock = block()
+  if (!resolvedBlock) return null
 
-                    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}
-                        />
+  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)
+                }
+
+                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>
-                    )
-                  })()}
-                </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>
+                      <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>
   )
 }