fix(mothership): keep a late write behind what the chat's queue already held

When the new-chat queue moved into a chat queue that already had messages,
migrate put those first, but a late write still used its index in the
new-chat queue and could land ahead of them. The move now records how many
messages it went behind, and late writes resolve their position, not just
their key (liveQueuePosition).

Author: waleedlatif1

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

@@ -136,6 +136,7 @@ import { useChatPanelStore } from '@/stores/chat-panel/store'
 import { useMothershipEffortStore } from '@/stores/mothership-effort/store'
 import {
   liveQueueKey,
+  liveQueuePosition,
   reusedRequestId,
   useMothershipQueueStore,
 } from '@/stores/mothership-queue/store'
@@ -4364,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 requeueKey = liveQueueKey(activeChatKey)
+      const { chatKey: requeueKey, index: requeueIndex } = liveQueuePosition(activeChatKey, 0)
       const chatless = requeueKey.startsWith(PENDING_CHAT_KEY_PREFIX)
       if (result.reason === 'withdrawn' && chatless) {
         handOffWithdrawnSend({ ...payload, userMessageId: result.userMessageId })
@@ -4374,7 +4375,7 @@ export function useChat(
          it, so anything queued while its POST was out was written after it. The one
          exception is a held send adopted from a dead mount of this surface in that
          window, which can be older; it lands behind this one. */
-      useMothershipQueueStore.getState().insertAt(requeueKey, 0, {
+      useMothershipQueueStore.getState().insertAt(requeueKey, requeueIndex, {
         ...createQueuedMessage(payload, result.userMessageId),
         ...requeuedFields(result.reason, 0, chatless ? heldSendSurface : undefined),
         admissionUnknown: result.admissionUnknown,
@@ -4933,7 +4934,10 @@ export function useChat(
         const withdrawnUserMessageId = withdrawn?.userMessageId
         /* The send may have waited on a Stop that saw the new chat's first message
            admitted, which moved this queue to that chat. */
-        const restoreKey = liveQueueKey(dispatchChatKey)
+        const { chatKey: restoreKey, index: restoreIndex } = liveQueuePosition(
+          dispatchChatKey,
+          originalIndex
+        )
         const chatless = restoreKey.startsWith(PENDING_CHAT_KEY_PREFIX)
         const savedHandoff = readQueuedSendHandoffState()
         const retainedHandoff =
@@ -4981,7 +4985,7 @@ export function useChat(
         }
         /** Once restored, the queue owns recovery; a second handoff reader must not resend it. */
         clearQueuedSendHandoffState(msg.id)
-        useMothershipQueueStore.getState().insertAt(restoreKey, originalIndex, {
+        useMothershipQueueStore.getState().insertAt(restoreKey, restoreIndex, {
           /* Only this outcome's policy applies: what an earlier one set (a hold, a
              retry delay, a surface) must not outlive it. */
           ...withoutRequeueFields(dispatched),

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

@@ -1,5 +1,9 @@
 import { beforeEach, describe, expect, it } from 'vitest'
-import { liveQueueKey, useMothershipQueueStore } from '@/stores/mothership-queue/store'
+import {
+  liveQueueKey,
+  liveQueuePosition,
+  useMothershipQueueStore,
+} from '@/stores/mothership-queue/store'
 import type { QueuedMothershipMessage } from '@/stores/mothership-queue/types'
 
 const message = (id: string, content = `content-${id}`): QueuedMothershipMessage => ({
@@ -190,6 +194,24 @@ describe('useMothershipQueueStore', () => {
       expect(liveQueueKey('chat-X')).toBe('chat-X')
     })
 
+    it('keeps a late write behind the messages the chat queue already held', () => {
+      useMothershipQueueStore.getState().enqueue('chat-Y', message('older-1'))
+      useMothershipQueueStore.getState().enqueue('chat-Y', message('older-2'))
+      useMothershipQueueStore.getState().enqueue('pending::moved', message('moved'))
+      useMothershipQueueStore.getState().migrate('pending::moved', 'chat-Y')
+
+      const position = liveQueuePosition('pending::moved', 0)
+      useMothershipQueueStore.getState().insertAt(position.chatKey, position.index, message('late'))
+
+      expect(position).toEqual({ chatKey: 'chat-Y', index: 2 })
+      expect(useMothershipQueueStore.getState().queues['chat-Y']?.map((m) => m.id)).toEqual([
+        'older-1',
+        'older-2',
+        'late',
+        'moved',
+      ])
+    })
+
     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

@@ -3,7 +3,11 @@ import { toError } from '@sim/utils/errors'
 import { toRecord, toRecordOrNull } from '@sim/utils/object'
 import { create } from 'zustand'
 import { createJSONStorage, devtools, persist } from 'zustand/middleware'
-import type { MothershipQueueState, QueuedMothershipMessage } from '@/stores/mothership-queue/types'
+import type {
+  MothershipQueueState,
+  QueuedMothershipMessage,
+  QueueMigration,
+} from '@/stores/mothership-queue/types'
 
 const logger = createLogger('MothershipQueueStore')
 
@@ -49,7 +53,7 @@ const initialState = {
   queues: {} as Record<string, QueuedMothershipMessage[]>,
   editing: {} as Record<string, string>,
   cleared: {} as Record<string, number>,
-  migratedTo: {} as Record<string, string>,
+  migratedTo: {} as Record<string, QueueMigration>,
 }
 
 /**
@@ -108,14 +112,27 @@ const setQueueForChat = (
   next.length === 0 ? omitKey(queues, chatKey) : { ...queues, [chatKey]: next }
 
 /**
- * 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.
+ * 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.
  */
-export function liveQueueKey(chatKey: string): string {
+export function liveQueuePosition(
+  chatKey: string,
+  index: number
+): { chatKey: string; index: number } {
   const { migratedTo } = useMothershipQueueStore.getState()
-  let key = chatKey
-  for (let hops = 0; hops < 8 && migratedTo[key] !== undefined; hops++) key = migratedTo[key]
-  return key
+  let position = { chatKey, index }
+  for (let hops = 0; hops < 8; hops++) {
+    const migration = migratedTo[position.chatKey]
+    if (!migration) break
+    position = { chatKey: migration.key, index: position.index + migration.behind }
+  }
+  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
 }
 
 export const useMothershipQueueStore = create<MothershipQueueState>()(
@@ -203,7 +220,10 @@ export const useMothershipQueueStore = create<MothershipQueueState>()(
         migrate: (fromKey, toKey) =>
           set((state) => {
             if (fromKey === toKey) return state
-            const migratedTo = { ...state.migratedTo, [fromKey]: toKey }
+            const migratedTo = {
+              ...state.migratedTo,
+              [fromKey]: { key: toKey, behind: state.queues[toKey]?.length ?? 0 },
+            }
             const fromQueue = state.queues[fromKey]
             const fromEditing = state.editing[fromKey]
             if (!fromQueue && fromEditing === undefined) return { migratedTo }

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

@@ -56,6 +56,13 @@ export type QueuedMessageEditPatch = Pick<
   | 'assistantSearchLevel'
 >
 
+/** 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
+}
+
 export interface MothershipQueueState {
   queues: Record<string, QueuedMothershipMessage[]>
   editing: Record<string, string>
@@ -68,9 +75,10 @@ export interface MothershipQueueState {
   /**
    * Where each new-chat key's queue moved when its chat became known
    * (`migrate`). A write that captured the old key before an `await` follows
-   * this (`liveQueueKey`), so it lands in the chat's queue, not a dead key.
+   * this (`liveQueueKey`, `liveQueuePosition`), so it lands in the chat's queue,
+   * not a dead key.
    */
-  migratedTo: Record<string, string>
+  migratedTo: Record<string, QueueMigration>
 
   enqueue: (chatKey: string, message: QueuedMothershipMessage) => void
   insertAt: (chatKey: string, index: number, message: QueuedMothershipMessage) => void