fix(mothership): anchor a late re-queue on the messages ahead of it, not an index

A restored message went back at an index captured at dispatch, offset by
how many messages a chat's queue held when the new-chat queue moved into it.
Removing any of those while the POST was out shifted it behind a newer
message. It now goes right after the last message still queued that was
ahead of it (at dispatch, or in the chat's queue before the move), else at
the head.

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts

@@ -4365,7 +4365,7 @@ export function useChat(
          chatless surface, whose key dies with the mount, goes to the
          cross-surface lanes. */
       /** The new-chat queue may have moved to its chat while the POST was out. */
-      const { chatKey: requeueKey, index: requeueIndex } = liveQueuePosition(activeChatKey, 0)
+      const { chatKey: requeueKey, index: requeueIndex } = liveQueuePosition(activeChatKey, [])
       const chatless = requeueKey.startsWith(PENDING_CHAT_KEY_PREFIX)
       if (result.reason === 'withdrawn' && chatless) {
         handOffWithdrawnSend({ ...payload, userMessageId: result.userMessageId })
@@ -4906,8 +4906,10 @@ export function useChat(
       const dispatchChatKey = chatKeyRef.current
       const queueAtStart =
         useMothershipQueueStore.getState().queues[dispatchChatKey] ?? EMPTY_MESSAGE_QUEUE
-      let originalIndex = queueAtStart.findIndex((queued) => queued.id === msg.id)
-      if (originalIndex === -1) {
+      const startIndex = queueAtStart.findIndex((queued) => queued.id === msg.id)
+      /** What was queued ahead of it, which it goes back behind if it is restored. */
+      let aheadIds = queueAtStart.slice(0, Math.max(0, startIndex)).map((queued) => queued.id)
+      if (startIndex === -1) {
         queuedMessageDispatchIds.delete(msg.id)
         return
       }
@@ -4936,7 +4938,7 @@ export function useChat(
            admitted, which moved this queue to that chat. */
         const { chatKey: restoreKey, index: restoreIndex } = liveQueuePosition(
           dispatchChatKey,
-          originalIndex
+          aheadIds
         )
         const chatless = restoreKey.startsWith(PENDING_CHAT_KEY_PREFIX)
         const savedHandoff = readQueuedSendHandoffState()
@@ -5009,7 +5011,7 @@ export function useChat(
         if (currentIndex === -1) {
           return
         }
-        originalIndex = currentIndex
+        aheadIds = queueAtSend.slice(0, currentIndex).map((queued) => queued.id)
 
         // Re-read live: the user may have applied an in-place edit (`replaceAt`)
         // between dispatch scheduling and this send.

apps/sim/stores/mothership-queue/store.test.ts

@@ -200,7 +200,7 @@ describe('useMothershipQueueStore', () => {
       useMothershipQueueStore.getState().enqueue('pending::moved', message('moved'))
       useMothershipQueueStore.getState().migrate('pending::moved', 'chat-Y')
 
-      const position = liveQueuePosition('pending::moved', 0)
+      const position = liveQueuePosition('pending::moved', [])
       useMothershipQueueStore.getState().insertAt(position.chatKey, position.index, message('late'))
 
       expect(position).toEqual({ chatKey: 'chat-Y', index: 2 })
@@ -212,6 +212,23 @@ describe('useMothershipQueueStore', () => {
       ])
     })
 
+    it('keeps a late write in order when messages ahead of it were removed meanwhile', () => {
+      useMothershipQueueStore.getState().enqueue('chat-Z', message('older'))
+      useMothershipQueueStore.getState().enqueue('pending::sent', message('later'))
+      useMothershipQueueStore.getState().migrate('pending::sent', 'chat-Z')
+      useMothershipQueueStore.getState().remove('chat-Z', 'older')
+
+      const position = liveQueuePosition('pending::sent', [])
+      useMothershipQueueStore
+        .getState()
+        .insertAt(position.chatKey, position.index, message('follow-up'))
+
+      expect(useMothershipQueueStore.getState().queues['chat-Z']?.map((m) => m.id)).toEqual([
+        'follow-up',
+        'later',
+      ])
+    })
+
     it('merges into an existing destination bucket instead of overwriting', () => {
       useMothershipQueueStore.getState().enqueue('chat-X', message('existing-1'))
       useMothershipQueueStore.getState().enqueue('chat-X', message('existing-2'))

apps/sim/stores/mothership-queue/store.ts

@@ -112,27 +112,42 @@ const setQueueForChat = (
   next.length === 0 ? omitKey(queues, chatKey) : { ...queues, [chatKey]: next }
 
 /**
- * Where a position a write captured before an `await` lies now: in the queue a
- * new-chat queue migrated to once its chat became known, if it did, behind the
- * messages that queue already held.
+ * The queue key a write captured before an `await` should use now: the key a
+ * new-chat queue migrated to once its chat became known, if it did.
+ */
+export function liveQueueKey(chatKey: string): string {
+  const { migratedTo } = useMothershipQueueStore.getState()
+  let key = chatKey
+  for (let hops = 0; hops < 8 && migratedTo[key] !== undefined; hops++) key = migratedTo[key].key
+  return key
+}
+
+/**
+ * Where a message goes back into its queue after a write captured before an
+ * `await`: in the queue's live key, right after the last message still there
+ * that was ahead of it (`aheadIds`, plus whatever a chat's queue already held
+ * when a new-chat queue moved into it), else at the head. Anchoring on ids, not
+ * an index, keeps it in order however the queue changed meanwhile.
  */
 export function liveQueuePosition(
   chatKey: string,
-  index: number
+  aheadIds: readonly string[]
 ): { chatKey: string; index: number } {
-  const { migratedTo } = useMothershipQueueStore.getState()
-  let position = { chatKey, index }
+  const { migratedTo, queues } = useMothershipQueueStore.getState()
+  const ahead = new Set(aheadIds)
+  let key = chatKey
   for (let hops = 0; hops < 8; hops++) {
-    const migration = migratedTo[position.chatKey]
+    const migration = migratedTo[key]
     if (!migration) break
-    position = { chatKey: migration.key, index: position.index + migration.behind }
+    for (const id of migration.ahead) ahead.add(id)
+    key = migration.key
   }
-  return position
-}
-
-/** The queue key a write captured before an `await` should use now (see `liveQueuePosition`). */
-export function liveQueueKey(chatKey: string): string {
-  return liveQueuePosition(chatKey, 0).chatKey
+  const queue = queues[key] ?? []
+  let index = 0
+  queue.forEach((message, position) => {
+    if (ahead.has(message.id)) index = position + 1
+  })
+  return { chatKey: key, index }
 }
 
 export const useMothershipQueueStore = create<MothershipQueueState>()(
@@ -222,7 +237,10 @@ export const useMothershipQueueStore = create<MothershipQueueState>()(
             if (fromKey === toKey) return state
             const migratedTo = {
               ...state.migratedTo,
-              [fromKey]: { key: toKey, behind: state.queues[toKey]?.length ?? 0 },
+              [fromKey]: {
+                key: toKey,
+                ahead: (state.queues[toKey] ?? []).map((message) => message.id),
+              },
             }
             const fromQueue = state.queues[fromKey]
             const fromEditing = state.editing[fromKey]

apps/sim/stores/mothership-queue/types.ts

@@ -59,8 +59,8 @@ export type QueuedMessageEditPatch = Pick<
 /** A new-chat queue's move to its chat's key. */
 export interface QueueMigration {
   key: string
-  /** Messages the chat's queue already held, which stay ahead of the moved ones. */
-  behind: number
+  /** Ids of the messages the chat's queue already held, which stay ahead of the moved ones. */
+  ahead: string[]
 }
 
 export interface MothershipQueueState {