From f0ed98222a3bf7da4d21a9bce94f78b0ccb903d0 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Sun, 30 Nov 2025 10:43:10 +0000
Subject: [PATCH] Skip reverted messages from display caches

---
 .../ui/src/components/message-stream-v2.tsx   | 23 ++++++++++++++-----
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/packages/ui/src/components/message-stream-v2.tsx b/packages/ui/src/components/message-stream-v2.tsx
index f65eabb8..30210c8b 100644
--- a/packages/ui/src/components/message-stream-v2.tsx
+++ b/packages/ui/src/components/message-stream-v2.tsx
@@ -192,8 +192,17 @@ 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(() =>
-    messageIds()
+    visibleMessageIds()
       .map((id) => store().getMessage(id))
       .filter((record): record is MessageRecord => Boolean(record)),
   )
@@ -250,15 +259,17 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
 
   const messageIndexMap = createMemo(() => {
     const map = new Map<string, number>()
-    const records = messageRecords()
-    records.forEach((record, index) => map.set(record.id, index))
+    const ids = visibleMessageIds()
+    ids.forEach((id, index) => map.set(id, index))
     return map
   })
 
   const lastAssistantIndex = createMemo(() => {
-    const records = messageRecords()
-    for (let index = records.length - 1; index >= 0; index--) {
-      if (records[index].role === "assistant") {
+    const ids = visibleMessageIds()
+    const resolvedStore = store()
+    for (let index = ids.length - 1; index >= 0; index--) {
+      const record = resolvedStore.getMessage(ids[index])
+      if (record?.role === "assistant") {
         return index
       }
     }