fix(tui): preserve initial user message on new session hydration

packages/tui/src/component/prompt/index.tsx

@@ -1091,6 +1091,17 @@ export function Prompt(props: PromptProps) {
       })
     } else {
       move.startSubmit()
+      sync.session.optimisticUserMessage({
+        sessionID,
+        text: inputText,
+        parts: [...editorParts, ...nonTextParts],
+        agent: agent.name,
+        model: {
+          providerID: selectedModel.providerID,
+          modelID: selectedModel.modelID,
+          variant,
+        },
+      })
       sdk.client.session
         .prompt(
           {
@@ -1111,6 +1122,7 @@ export function Prompt(props: PromptProps) {
           { throwOnError: true },
         )
         .catch((error) => {
+          sync.session.clearOptimisticUserMessage(sessionID)
           toast.show({
             title: "Failed to send prompt",
             message: errorMessage(error),
@@ -1131,15 +1143,14 @@ export function Prompt(props: PromptProps) {
     setStore("extmarkToPartIndex", new Map())
     props.onSubmit?.()
 
-    // temporary hack to make sure the message is sent
     if (!props.sessionID) {
       if (editorParts.length > 0) editor.preserveSelectionFromNewSession()
-      setTimeout(() => {
+      queueMicrotask(() => {
         route.navigate({
           type: "session",
           sessionID,
         })
-      }, 50)
+      })
     }
     input.clear()
     if (finishMoveProgress) move.finishSubmit()

packages/tui/src/context/sync.tsx

@@ -144,6 +144,32 @@ export const {
     const fullSyncedSessions = new Set<string>()
     const syncingSessions = new Map<string, Promise<void>>()
     const hydratingSessions = new Map<string, { messages: Set<string>; parts: Set<string> }>()
+    const optimisticBySession = new Map<string, string>()
+
+    function clearOptimisticUserMessage(sessionID: string) {
+      const messageID = optimisticBySession.get(sessionID)
+      if (!messageID) return
+      optimisticBySession.delete(sessionID)
+      const messages = store.message[sessionID]
+      if (messages) {
+        const match = search(messages, messageID, (m) => m.id)
+        if (match.found) {
+          setStore(
+            "message",
+            sessionID,
+            produce((draft) => {
+              draft.splice(match.index, 1)
+            }),
+          )
+        }
+      }
+      setStore(
+        "part",
+        produce((draft) => {
+          delete draft[messageID]
+        }),
+      )
+    }
     const touchMessage = (sessionID: string, messageID: string) => {
       hydratingSessions.get(sessionID)?.messages.add(messageID)
     }
@@ -313,6 +339,12 @@ export const {
         }
 
         case "message.updated": {
+          if (event.properties.info.role === "user") {
+            const optimisticID = optimisticBySession.get(event.properties.info.sessionID)
+            if (optimisticID && optimisticID !== event.properties.info.id) {
+              clearOptimisticUserMessage(event.properties.info.sessionID)
+            }
+          }
           touchMessage(event.properties.info.sessionID, event.properties.info.id)
           const messages = store.message[event.properties.info.sessionID]
           if (!messages) {
@@ -605,14 +637,33 @@ export const {
                 if (!match.found) draft.session.splice(match.index, 0, session.data!)
                 draft.todo[sessionID] = todo.data ?? []
                 const currentMessages = draft.message[sessionID] ?? []
+                const optimisticID = optimisticBySession.get(sessionID)
+                // Once the server has persisted the real first user message, drop the
+                // optimistic placeholder — otherwise both the real and the optimistic
+                // user message render as duplicates. clearOptimisticUserMessage's message
+                // removal is handled by the `visible` replacement below; here we only need
+                // to forget the id and clean up its parts.
+                const optimisticSuperseded =
+                  optimisticID !== undefined &&
+                  (messages.data ?? []).some(
+                    (message) => message.info.role === "user" && message.info.id !== optimisticID,
+                  )
+                if (optimisticID !== undefined && optimisticSuperseded) {
+                  optimisticBySession.delete(sessionID)
+                  delete draft.part[optimisticID]
+                }
+                const optimisticIDs =
+                  optimisticID && !optimisticSuperseded ? new Set([optimisticID]) : new Set<string>()
                 const infos = (messages.data ?? []).flatMap((message) => {
                   if (!tracker.messages.has(message.info.id)) return [message.info]
                   const current = currentMessages.find((item) => item.id === message.info.id)
                   return current ? [current] : []
                 })
                 infos.push(
                   ...currentMessages.filter(
-                    (message) => tracker.messages.has(message.id) && !infos.some((item) => item.id === message.id),
+                    (message) =>
+                      (tracker.messages.has(message.id) || optimisticIDs.has(message.id)) &&
+                      !infos.some((item) => item.id === message.id),
                   ),
                 )
                 const removed = infos.slice(0, -100)
@@ -658,6 +709,63 @@ export const {
           syncingSessions.set(sessionID, task)
           return task
         },
+        optimisticUserMessage(input: {
+          sessionID: string
+          text: string
+          parts?: Array<Omit<Part, "id" | "sessionID" | "messageID">>
+          agent: string
+          model: {
+            providerID: string
+            modelID: string
+            variant?: string
+          }
+        }) {
+          clearOptimisticUserMessage(input.sessionID)
+          const messageID = `opt_${crypto.randomUUID()}`
+          optimisticBySession.set(input.sessionID, messageID)
+          const created = Date.now()
+          const message: Message = {
+            id: messageID,
+            sessionID: input.sessionID,
+            role: "user",
+            time: { created },
+            agent: input.agent,
+            model: input.model,
+          }
+          const parts: Part[] = [
+            ...(input.parts ?? []).map(
+              (part) =>
+                ({
+                  ...part,
+                  id: `opt_${crypto.randomUUID()}`,
+                  sessionID: input.sessionID,
+                  messageID,
+                }) as Part,
+            ),
+            {
+              id: `opt_${crypto.randomUUID()}`,
+              sessionID: input.sessionID,
+              messageID,
+              type: "text",
+              text: input.text,
+            },
+          ]
+          const messages = store.message[input.sessionID]
+          if (!messages) {
+            setStore("message", input.sessionID, [message])
+          } else {
+            setStore(
+              "message",
+              input.sessionID,
+              produce((draft) => {
+                draft.push(message)
+              }),
+            )
+          }
+          setStore("part", messageID, parts)
+          return messageID
+        },
+        clearOptimisticUserMessage,
       },
       bootstrap,
     }

packages/tui/test/cli/cmd/tui/sync-live-hydration.test.tsx

@@ -260,3 +260,130 @@ test("a message removed during hydration does not regain stale parts", async ()
     app.renderer.destroy()
   }
 })
+
+test("optimistic user message survives new session hydration race", async () => {
+  await using tmp = await tmpdir()
+  await Bun.write(`${tmp.path}/kv.json`, "{}")
+
+  const userText = "What is Rust?"
+  let resolveMessages!: (response: Response) => void
+  const messages = new Promise<Response>((resolve) => {
+    resolveMessages = resolve
+  })
+  let requested = false
+  const { app, emit, sync } = await mount((url) => {
+    if (url.pathname === `/session/${sessionID}`) return json(session)
+    if (url.pathname === `/session/${sessionID}/message`) {
+      requested = true
+      return messages
+    }
+    if (url.pathname === `/session/${sessionID}/todo` || url.pathname === `/session/${sessionID}/diff`) return json([])
+    return undefined
+  }, tmp.path)
+
+  try {
+    sync.session.optimisticUserMessage({
+      sessionID,
+      text: userText,
+      agent: "build",
+      model: { providerID: "test", modelID: "model" },
+    })
+
+    const hydrate = sync.session.sync(sessionID)
+    await wait(() => requested)
+    emit(global({ id: "evt_message", type: "message.updated", properties: { sessionID, info: assistant } }))
+    emit(
+      global({
+        id: "evt_part",
+        type: "message.part.updated",
+        properties: {
+          sessionID,
+          time: 2,
+          part: { id: partID, sessionID, messageID, type: "text", text: "Hi! How can I help you?" },
+        },
+      }),
+    )
+    await wait(() => sync.data.part[messageID]?.[0]?.type === "text")
+
+    resolveMessages(
+      json([
+        {
+          info: assistant,
+          parts: [{ id: partID, sessionID, messageID, type: "text", text: "Hi! How can I help you?" }],
+        },
+      ]),
+    )
+    await hydrate
+
+    const userMessages = (sync.data.message[sessionID] ?? []).filter((message) => message.role === "user")
+    expect(userMessages).toHaveLength(1)
+    const userParts = sync.data.part[userMessages[0]!.id] ?? []
+    expect(userParts.some((part) => part.type === "text" && part.text === userText)).toBe(true)
+    expect(sync.data.part[messageID][0]).toMatchObject({ text: "Hi! How can I help you?" })
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
+test("hydration returning the real user message drops the optimistic duplicate", async () => {
+  await using tmp = await tmpdir()
+  await Bun.write(`${tmp.path}/kv.json`, "{}")
+
+  const userText = "What is Rust?"
+  const realUserID = "msg_user"
+  const realUserPartID = "prt_user"
+  const realUser = {
+    id: realUserID,
+    sessionID,
+    role: "user" as const,
+    agent: "build",
+    model: { providerID: "test", modelID: "model" },
+    time: { created: 1 },
+  }
+  let resolveMessages!: (response: Response) => void
+  const messages = new Promise<Response>((resolve) => {
+    resolveMessages = resolve
+  })
+  let requested = false
+  const { app, sync } = await mount((url) => {
+    if (url.pathname === `/session/${sessionID}`) return json(session)
+    if (url.pathname === `/session/${sessionID}/message`) {
+      requested = true
+      return messages
+    }
+    if (url.pathname === `/session/${sessionID}/todo` || url.pathname === `/session/${sessionID}/diff`) return json([])
+    return undefined
+  }, tmp.path)
+
+  try {
+    sync.session.optimisticUserMessage({
+      sessionID,
+      text: userText,
+      agent: "build",
+      model: { providerID: "test", modelID: "model" },
+    })
+    const optimisticUserID = sync.data.message[sessionID]!.find((message) => message.role === "user")!.id
+    expect(optimisticUserID.startsWith("opt_")).toBe(true)
+
+    const hydrate = sync.session.sync(sessionID)
+    await wait(() => requested)
+
+    // Server has now persisted the real user message; the hydration response
+    // carries both it and the assistant reply.
+    resolveMessages(
+      json([
+        { info: realUser, parts: [{ id: realUserPartID, sessionID, messageID: realUserID, type: "text", text: userText }] },
+        { info: assistant, parts: [{ id: partID, sessionID, messageID, type: "text", text: "Hi!" }] },
+      ]),
+    )
+    await hydrate
+
+    const userMessages = (sync.data.message[sessionID] ?? []).filter((message) => message.role === "user")
+    expect(userMessages).toHaveLength(1)
+    expect(userMessages[0]!.id).toBe(realUserID)
+    // The optimistic placeholder and its parts must be gone.
+    expect(sync.data.part[optimisticUserID]).toBeUndefined()
+  } finally {
+    app.renderer.destroy()
+  }
+})