From cd46f0e3fb2ac3e8a064adf5a236fb87ae489274 Mon Sep 17 00:00:00 2001
From: Bohdan Triapitsyn <artmore@protonmail.com>
Date: Sat, 22 Aug 2026 01:55:52 +0300
Subject: [PATCH] fix(chat): stop the just-sent user message from flickering

The server echo of an optimistic user message carries different part
ids, so the reducer dropped the optimistic part and appended the server
one at the end, and MessageBody keyed user parts by part id. The key
change remounted the text subtree (blank frame, markdown re-parse,
truncation state reset) and the append reordered text against file
parts.

Replace the optimistic part in place and key user parts positionally.
---
 .../src/components/chat/message/MessageBody.tsx  |  9 ++++++---
 .../ui/src/sync/__tests__/event-reducer.test.ts  | 16 ++++++++++++++++
 packages/ui/src/sync/event-reducer.ts            |  7 +++++--
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/packages/ui/src/components/chat/message/MessageBody.tsx b/packages/ui/src/components/chat/message/MessageBody.tsx
index fa0bbfa3..c6c09f46 100644
--- a/packages/ui/src/components/chat/message/MessageBody.tsx
+++ b/packages/ui/src/components/chat/message/MessageBody.tsx
@@ -726,10 +726,13 @@ const UserMessageBody = React.memo(({ messageId, parts, messageCreatedAt, isMobi
                 )}
                 style={useStickyScrollableUserContent ? { maxHeight: 'calc(var(--chat-scroll-height, 100dvh) * 0.4)' } : undefined}
             >
+                {/* Positional keys, not part ids: the server echo of a just-sent
+                    message swaps the optimistic part id, and id-based keys would
+                    remount the text subtree (blank frame + height jump). */}
                 {userContentParts.map((part, index) => {
                     if (isSubtaskPart(part)) {
                         return (
-                            <React.Fragment key={part.id ?? `user-subtask-${index}`}>
+                            <React.Fragment key={`user-subtask-${index}`}>
                                 <UserSubtaskPart part={part} />
                             </React.Fragment>
                         );
@@ -737,7 +740,7 @@ const UserMessageBody = React.memo(({ messageId, parts, messageCreatedAt, isMobi
 
                     if (isShellActionPart(part)) {
                         return (
-                            <React.Fragment key={part.id ?? `user-shell-${index}`}>
+                            <React.Fragment key={`user-shell-${index}`}>
                                 <UserShellActionPart part={part} />
                             </React.Fragment>
                         );
@@ -752,7 +755,7 @@ const UserMessageBody = React.memo(({ messageId, parts, messageCreatedAt, isMobi
                         }
                     }
                     return (
-                        <React.Fragment key={part.id ?? `user-text-${index}`}>
+                        <React.Fragment key={`user-text-${index}`}>
                             <UserTextPart
                                 part={part}
                                 messageId={messageId}
diff --git a/packages/ui/src/sync/__tests__/event-reducer.test.ts b/packages/ui/src/sync/__tests__/event-reducer.test.ts
index 59ffc74a..2dda4914 100644
--- a/packages/ui/src/sync/__tests__/event-reducer.test.ts
+++ b/packages/ui/src/sync/__tests__/event-reducer.test.ts
@@ -114,6 +114,22 @@ describe("applyDirectoryEvent", () => {
     expect(draft.part.msg_1).toEqual([legacyPart, currentPart])
   })
 
+  test("replaces an optimistic user part in place instead of appending it", () => {
+    const optimisticText = { id: "prt_optimistic_text", messageID: "msg_1", type: "text", text: "hi" } as Part
+    const optimisticFile = { id: "prt_optimistic_file", messageID: "msg_1", type: "file", filename: "a.png" } as Part
+    const serverText = { id: "prt_server_text", messageID: "msg_1", sessionID: "ses_1", type: "text", text: "hi" } as Part
+    const draft = state({
+      message: { ses_1: [{ id: "msg_1", sessionID: "ses_1", role: "user", time: { created: 1 } } as Message] },
+      part: { msg_1: [optimisticText, optimisticFile] },
+    })
+
+    expect(applyDirectoryEvent(draft, {
+      type: "message.part.updated",
+      properties: { part: serverText },
+    } as Event)).toBe(true)
+    expect(draft.part.msg_1).toEqual([serverText, optimisticFile])
+  })
+
   test("returns typed materialization when delta arrives before parts", () => {
     const result = applyDirectoryEvent(state(), deltaEvent())
 
diff --git a/packages/ui/src/sync/event-reducer.ts b/packages/ui/src/sync/event-reducer.ts
index 0eaf4e71..13057258 100644
--- a/packages/ui/src/sync/event-reducer.ts
+++ b/packages/ui/src/sync/event-reducer.ts
@@ -441,9 +441,12 @@ export function applyDirectoryEvent(
           ? next.findIndex((p) => p.type === part.type && !(p as { sessionID?: string }).sessionID)
           : -1
         if (optimisticIndex >= 0) {
-          next.splice(optimisticIndex, 1)
+          // Replace in place: pushing to the end reorders text/file parts of a
+          // just-sent message and remounts its rendered subtree.
+          next[optimisticIndex] = part
+        } else {
+          next.push(part)
         }
-        next.push(part)
       }
       draft.part[messageID] = next
       return missingOwningMessage