From 57c16052427f8f39bd6bcc6fbdcc74eb203c0478 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Wed, 3 Dec 2025 17:07:05 +0000
Subject: [PATCH] Message addition performance improvements

---
 .../ui/src/stores/message-v2/instance-store.ts  | 15 ---------------
 packages/ui/src/stores/session-events.ts        | 17 +++++++----------
 2 files changed, 7 insertions(+), 25 deletions(-)

diff --git a/packages/ui/src/stores/message-v2/instance-store.ts b/packages/ui/src/stores/message-v2/instance-store.ts
index 19509be8..693af550 100644
--- a/packages/ui/src/stores/message-v2/instance-store.ts
+++ b/packages/ui/src/stores/message-v2/instance-store.ts
@@ -296,9 +296,6 @@ export function createInstanceMessageStore(instanceId: string): InstanceMessageS
     ensureSessionEntry(sessionId)
 
     const incomingIds = inputs.map((item) => item.id)
-    const incomingIdSet = new Set(incomingIds)
-    const existingIds = state.sessions[sessionId]?.messageIds ?? []
-    const removedIds = existingIds.filter((id) => !incomingIdSet.has(id))
 
     const normalizedRecords: Record<string, MessageRecord> = {}
     const now = Date.now()
@@ -331,18 +328,6 @@ export function createInstanceMessageStore(instanceId: string): InstanceMessageS
       ...state.permissions.byMessage,
     }
 
-    removedIds.forEach((id) => {
-      if (nextMessages[id]?.sessionId === sessionId) {
-        delete nextMessages[id]
-        delete nextMessageInfoVersion[id]
-        delete nextPendingParts[id]
-        if (nextPermissionsByMessage[id]) {
-          delete nextPermissionsByMessage[id]
-        }
-      }
-      messageInfoCache.delete(id)
-    })
-
     Object.entries(normalizedRecords).forEach(([id, record]) => {
       nextMessages[id] = record
     })
diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts
index 5ae404f8..2a0cfd5c 100644
--- a/packages/ui/src/stores/session-events.ts
+++ b/packages/ui/src/stores/session-events.ts
@@ -62,16 +62,13 @@ function findPendingMessageId(
   role: MessageRole,
 ): string | undefined {
   const messageIds = store.getSessionMessageIds(sessionId)
-  for (let i = messageIds.length - 1; i >= 0; i -= 1) {
-    const record = store.getMessage(messageIds[i])
-    if (!record) continue
-    if (record.sessionId !== sessionId) continue
-    if (record.role !== role) continue
-    if (record.status === "sending") {
-      return record.id
-    }
-  }
-  return undefined
+  const lastId = messageIds[messageIds.length - 1]
+  if (!lastId) return undefined
+  const record = store.getMessage(lastId)
+  if (!record) return undefined
+  if (record.sessionId !== sessionId) return undefined
+  if (record.role !== role) return undefined
+  return record.status === "sending" ? record.id : undefined
 }
 
 function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | MessagePartUpdatedEvent): void {