fix(tui): order messages by creation time (#40994)

Co-authored-by: Dax <mail@thdxr.com>

packages/tui/src/context/sync.tsx

@@ -51,6 +51,12 @@ function search<T>(items: T[], target: string, key: (item: T) => string) {
   return { found: false, index: left }
 }
 
+function compareMessage(a: Message, b: Message) {
+  return a.time.created - b.time.created || a.id.localeCompare(b.id)
+}
+
+const messageKey = (message: Message) => message.time.created + message.id
+
 export const {
   context: SyncContext,
   use: useSync,
@@ -319,7 +325,7 @@ export const {
             setStore("message", event.properties.info.sessionID, [event.properties.info])
             break
           }
-          const result = search(messages, event.properties.info.id, (m) => m.id)
+          const result = search(messages, messageKey(event.properties.info), messageKey)
           if (result.found) {
             setStore("message", event.properties.info.sessionID, result.index, reconcile(event.properties.info))
             break
@@ -355,13 +361,13 @@ export const {
         case "message.removed": {
           touchMessage(event.properties.sessionID, event.properties.messageID)
           const messages = store.message[event.properties.sessionID]
-          const result = search(messages, event.properties.messageID, (m) => m.id)
-          if (result.found) {
+          const index = messages.findIndex((message) => message.id === event.properties.messageID)
+          if (index !== -1) {
             setStore(
               "message",
               event.properties.sessionID,
               produce((draft) => {
-                draft.splice(result.index, 1)
+                draft.splice(index, 1)
               }),
             )
           }
@@ -374,7 +380,7 @@ export const {
             setStore("part", event.properties.part.messageID, [event.properties.part])
             break
           }
-          const result = search(parts, event.properties.part.id, (p) => p.id)
+          const result = search(parts, event.properties.part.id, (part) => part.id)
           if (result.found) {
             setStore("part", event.properties.part.messageID, result.index, reconcile(event.properties.part))
             break
@@ -392,7 +398,7 @@ export const {
         case "message.part.delta": {
           const parts = store.part[event.properties.messageID]
           if (!parts) break
-          const result = search(parts, event.properties.partID, (p) => p.id)
+          const result = search(parts, event.properties.partID, (part) => part.id)
           if (!result.found) break
           touchPart(event.properties.sessionID, event.properties.partID)
           setStore(
@@ -411,7 +417,7 @@ export const {
         case "message.part.removed": {
           touchPart(event.properties.sessionID, event.properties.partID)
           const parts = store.part[event.properties.messageID]
-          const result = search(parts, event.properties.partID, (p) => p.id)
+          const result = search(parts, event.properties.partID, (part) => part.id)
           if (result.found) {
             setStore(
               "part",
@@ -615,6 +621,7 @@ export const {
                     (message) => tracker.messages.has(message.id) && !infos.some((item) => item.id === message.id),
                   ),
                 )
+                infos.sort(compareMessage)
                 const removed = infos.slice(0, -100)
                 const visible = infos.slice(-100)
                 const visibleIDs = new Set(visible.map((message) => message.id))

packages/tui/src/routes/session/index.tsx

@@ -211,6 +211,12 @@ export function Session() {
       .toSorted((a, b) => (a.id < b.id ? -1 : a.id > b.id ? 1 : 0))
   })
   const messages = createMemo(() => sync.data.message[route.sessionID] ?? [])
+  const messagesBeforeRevert = () => {
+    const messageID = session()?.revert?.messageID
+    if (!messageID) return messages()
+    const index = messages().findIndex((message) => message.id === messageID)
+    return index === -1 ? messages() : messages().slice(0, index)
+  }
   const foregroundTasks = createMemo(() =>
     sync.data.capabilities.experimentalBackgroundSubagents
       ? messages().flatMap((message) =>
@@ -236,9 +242,11 @@ export function Session() {
   const disabled = createMemo(() => permissions().length > 0 || questions().length > 0)
 
   const pending = createMemo(() => {
-    const completed = messages().findLast((x) => x.role === "assistant" && x.time.completed)?.id
-    return messages().findLast((x) => x.role === "assistant" && !x.time.completed && (!completed || x.id > completed))
-      ?.id
+    const completed = messages().findLastIndex((message) => message.role === "assistant" && message.time.completed)
+    const pending = messages().findLastIndex(
+      (message, index) => index > completed && message.role === "assistant" && !message.time.completed,
+    )
+    return pending === -1 ? undefined : pending
   })
 
   const lastAssistant = createMemo(() => {
@@ -610,8 +618,7 @@ export function Session() {
       run: async () => {
         const status = sync.data.session_status?.[route.sessionID]
         if (status?.type !== "idle") await sdk.client.session.abort({ sessionID: route.sessionID }).catch(() => {})
-        const revert = session()?.revert?.messageID
-        const message = messages().findLast((x) => (!revert || x.id < revert) && x.role === "user")
+        const message = messagesBeforeRevert().findLast((item) => item.role === "user")
         if (!message) return
         void sdk.client.session
           .revert({
@@ -872,10 +879,7 @@ export function Session() {
       value: "messages.copy",
       category: "Session",
       run: () => {
-        const revertID = session()?.revert?.messageID
-        const lastAssistantMessage = messages().findLast(
-          (msg) => msg.role === "assistant" && (!revertID || msg.id < revertID),
-        )
+        const lastAssistantMessage = messagesBeforeRevert().findLast((message) => message.role === "assistant")
         if (!lastAssistantMessage) {
           toast.show({ message: "No assistant messages found", variant: "error" })
           dialog.clear()
@@ -1118,13 +1122,22 @@ export function Session() {
 
   const revertInfo = createMemo(() => session()?.revert)
   const revertMessageID = createMemo(() => revertInfo()?.messageID)
+  const revertMessageIndex = createMemo(() => {
+    const messageID = revertMessageID()
+    if (!messageID) return -1
+    return messages().findIndex((message) => message.id === messageID)
+  })
 
   const revertDiffFiles = createMemo(() => getRevertDiffFiles(revertInfo()?.diff ?? ""))
 
   const revertRevertedMessages = createMemo(() => {
     const messageID = revertMessageID()
     if (!messageID) return []
-    return messages().filter((x) => x.id >= messageID && x.role === "user")
+    const index = revertMessageIndex()
+    if (index === -1) return []
+    return messages()
+      .slice(index)
+      .filter((message) => message.role === "user")
   })
 
   const revert = createMemo(() => {
@@ -1247,7 +1260,9 @@ export function Session() {
                           )
                         })()}
                       </Match>
-                      <Match when={revert()?.messageID && message.id >= revert()!.messageID}>
+                      <Match
+                        when={revert()?.messageID && revertMessageIndex() !== -1 && index() >= revertMessageIndex()}
+                      >
                         <></>
                       </Match>
                       <Match when={message.role === "user"}>
@@ -1352,7 +1367,7 @@ function UserMessage(props: {
   parts: Part[]
   onMouseUp: () => void
   index: number
-  pending?: string
+  pending?: number
 }) {
   const ctx = use()
   const local = useLocal()
@@ -1370,7 +1385,7 @@ function UserMessage(props: {
   const files = createMemo(() => props.parts.flatMap((x) => (x.type === "file" ? [x] : [])))
   const { theme } = useTheme()
   const [hover, setHover] = createSignal(false)
-  const queued = createMemo(() => props.pending && props.message.id > props.pending)
+  const queued = createMemo(() => props.pending !== undefined && props.index > props.pending)
   const color = createMemo(() => local.agent.color(props.message.agent))
   const queuedFg = createMemo(() => selectedForeground(theme, color()))
   const metadataVisible = createMemo(() => queued() || ctx.showTimestamps())

packages/tui/src/util/transcript.ts

@@ -35,7 +35,9 @@ export function formatTranscript(
   transcript += `**Updated:** ${new Date(session.time.updated).toLocaleString()}\n\n`
   transcript += `---\n\n`
 
-  for (const msg of messages) {
+  for (const msg of messages.toSorted(
+    (a, b) => a.info.time.created - b.info.time.created || a.info.id.localeCompare(b.info.id),
+  )) {
     transcript += formatMessage(msg.info, msg.parts, options, providers)
     transcript += `---\n\n`
   }

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

@@ -33,6 +33,29 @@ function global(payload: GlobalEvent["payload"]): GlobalEvent {
   return { directory: "/tmp/other", project: "proj_test", payload }
 }
 
+test("live messages use creation time with an ID tie-break", async () => {
+  await using tmp = await tmpdir()
+  await Bun.write(`${tmp.path}/kv.json`, "{}")
+  const { app, emit, sync } = await mount(undefined, tmp.path)
+  const messages = [
+    { ...assistant, id: "msg_a", time: { created: 30, completed: 31 } },
+    { ...assistant, id: "msg_z", time: { created: 10, completed: 11 } },
+    { ...assistant, id: "msg_m", time: { created: 20, completed: 21 } },
+    { ...assistant, id: "msg_b", time: { created: 20, completed: 21 } },
+  ]
+
+  try {
+    for (const info of messages) {
+      emit(global({ id: `evt_${info.id}`, type: "message.updated", properties: { sessionID, info } }))
+    }
+    await wait(() => sync.data.message[sessionID]?.length === messages.length)
+
+    expect(sync.data.message[sessionID].map((message) => message.id)).toEqual(["msg_z", "msg_b", "msg_m", "msg_a"])
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("stale session hydration does not overwrite live message parts", async () => {
   await using tmp = await tmpdir()
   await Bun.write(`${tmp.path}/kv.json`, "{}")

packages/tui/test/util/transcript.test.ts

@@ -349,6 +349,34 @@ describe("transcript", () => {
       expect(result).toContain("---")
     })
 
+    test("orders messages by creation time and preserves part order", () => {
+      const message = (id: string, created: number, parts: string[]) => ({
+        info: {
+          id,
+          sessionID: "ses_abc123",
+          role: "user" as const,
+          agent: "build",
+          model: { providerID: "anthropic", modelID: "claude" },
+          time: { created },
+        },
+        parts: parts.map((text, index) => ({
+          id: `part_${parts.length - index}`,
+          sessionID: "ses_abc123",
+          messageID: id,
+          type: "text" as const,
+          text,
+        })),
+      })
+      const result = formatTranscript(
+        { id: "ses_abc123", title: "Order", time: { created: 1, updated: 2 } },
+        [message("msg_a", 30, ["third"]), message("msg_z", 10, ["first", "second"])],
+        { thinking: false, toolDetails: false, assistantMetadata: false },
+      )
+
+      expect(result.indexOf("first")).toBeLessThan(result.indexOf("second"))
+      expect(result.indexOf("second")).toBeLessThan(result.indexOf("third"))
+    })
+
     test("falls back to raw model id when provider data is missing", () => {
       const session = {
         id: "ses_abc123",