fix(mothership): roll back a failed effort save by pick, not by value (#8655)

Picks queued behind an in-flight save run onMutate at once, so after low, high, low a failed first save dropped the newest low and the picker fell back to the stale server effort. Each pick now carries a token and a failed save drops only its own pick.

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/home/components/user-input/components/model-selector.tsx

@@ -35,7 +35,7 @@ export function ModelSelector() {
   const { chatId } = useChatSurface()
   const { data: chatHistory } = useMothershipChatHistory(chatId)
   const chatPick = useMothershipEffortStore((state) =>
-    chatId ? state.chatEfforts[chatId] : undefined
+    chatId ? state.chatEfforts[chatId]?.effort : undefined
   )
   const newChatEffort = useMothershipEffortStore((state) => state.newChatEffort)
   const setNewChatEffort = useMothershipEffortStore((state) => state.setNewChatEffort)

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

@@ -3411,7 +3411,7 @@ export function useChat(
         options?.requestMode === 'assistant'
           ? undefined
           : requestChatId
-            ? (effortStore.chatEfforts[requestChatId] ??
+            ? (effortStore.chatEfforts[requestChatId]?.effort ??
               queryClient.getQueryData<MothershipChatHistory>(
                 mothershipChatKeys.detail(requestChatId)
               )?.effort)

apps/sim/hooks/queries/mothership-chats.test.ts

@@ -1,6 +1,7 @@
 import { jsonResponse } from '@sim/testing/helpers/http'
 import { reactQueryMock, reactQueryMockFns } from '@sim/testing/mocks/react-query.mock'
 import { sleep } from '@sim/utils/helpers'
+import type { MutationObserverOptions } from '@tanstack/react-query'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const { suspendBrowserScope, suspendTerminalScope, clearChat } = vi.hoisted(() => ({
@@ -23,7 +24,12 @@ vi.mock('@/lib/terminal/transport', () => ({
   suspendTerminalScope,
 }))
 
-import { useDeleteMothershipChats } from '@/hooks/queries/mothership-chats'
+import type { MothershipEffort } from '@/lib/mothership/model-options'
+import {
+  useDeleteMothershipChats,
+  useSetMothershipChatEffort,
+} from '@/hooks/queries/mothership-chats'
+import { useMothershipEffortStore } from '@/stores/mothership-effort/store'
 
 const queryClient = reactQueryMockFns.mockQueryClient
 
@@ -93,4 +99,40 @@ describe('tasks query boundary parsing', () => {
       queryKey: ['mothership-chats', 'detail', 'chat-b'],
     })
   })
+
+  it('keeps the latest effort pick when an earlier queued save of the same value fails', async () => {
+    const tanstack =
+      await vi.importActual<typeof import('@tanstack/react-query')>('@tanstack/react-query')
+    const client = new tanstack.QueryClient()
+    const observer = new tanstack.MutationObserver(
+      client,
+      useSetMothershipChatEffort('chat-1') as unknown as MutationObserverOptions<
+        void,
+        Error,
+        MothershipEffort,
+        { pick: number }
+      >
+    )
+    const saves = [
+      Promise.withResolvers<Response>(),
+      Promise.withResolvers<Response>(),
+      Promise.withResolvers<Response>(),
+    ]
+    for (const save of saves) vi.mocked(fetch).mockReturnValueOnce(save.promise)
+    useMothershipEffortStore.getState().reset()
+    const outcomes = (['low', 'high', 'low'] as const).map((effort) =>
+      observer.mutate(effort).catch(() => undefined)
+    )
+    await sleep(1)
+    expect(fetch).toHaveBeenCalledTimes(1)
+
+    saves[0].resolve(new Response('save failed', { status: 500 }))
+    await outcomes[0]
+    expect(useMothershipEffortStore.getState().chatEfforts['chat-1']?.effort).toBe('low')
+
+    saves[1].resolve(jsonResponse({ success: true }))
+    saves[2].resolve(jsonResponse({ success: true }))
+    await Promise.all(outcomes)
+    expect(useMothershipEffortStore.getState().chatEfforts['chat-1']?.effort).toBe('low')
+  })
 })

apps/sim/hooks/queries/mothership-chats.ts

@@ -610,11 +610,15 @@ export function useSetMothershipChatEffort(chatId: string | undefined) {
       return setChatEffort({ chatId, effort })
     },
     scope: { id: `mothership-chat-effort:${chatId ?? ''}` },
+    // Runs at once even while an earlier save for this chat holds the scope, so a failed
+    // save rolls back its own pick by token, never a later pick of the same value.
     onMutate: (effort) => {
-      if (chatId) useMothershipEffortStore.getState().setChatEffort(chatId, effort)
+      if (!chatId) return undefined
+      return { pick: useMothershipEffortStore.getState().setChatEffort(chatId, effort) }
     },
-    onError: (_error, effort) => {
-      if (chatId) useMothershipEffortStore.getState().dropChatEffort(chatId, effort)
+    onError: (_error, _effort, context) => {
+      if (chatId && context)
+        useMothershipEffortStore.getState().dropChatEffort(chatId, context.pick)
     },
     onSuccess: (_data, effort) => {
       queryClient.setQueryData<MothershipChatHistory>(

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

@@ -42,11 +42,30 @@ describe('Build reasoning preferences', () => {
 
   it('keeps a newer chat pick when an older pick fails to save', () => {
     const store = useMothershipEffortStore.getState()
-    store.setChatEffort('chat-1', 'low')
+    const low = store.setChatEffort('chat-1', 'low')
+    const high = store.setChatEffort('chat-1', 'high')
+    store.dropChatEffort('chat-1', low)
+    expect(useMothershipEffortStore.getState().chatEfforts['chat-1']?.effort).toBe('high')
+    store.dropChatEffort('chat-1', high)
+    expect(useMothershipEffortStore.getState().chatEfforts).toEqual({})
+  })
+
+  it('keeps a newer pick of the same value when the first pick fails to save', () => {
+    const store = useMothershipEffortStore.getState()
+    const first = store.setChatEffort('chat-1', 'low')
     store.setChatEffort('chat-1', 'high')
-    store.dropChatEffort('chat-1', 'low')
-    expect(useMothershipEffortStore.getState().chatEfforts).toEqual({ 'chat-1': 'high' })
-    store.dropChatEffort('chat-1', 'high')
+    const latest = store.setChatEffort('chat-1', 'low')
+    store.dropChatEffort('chat-1', first)
+    expect(useMothershipEffortStore.getState().chatEfforts['chat-1']?.effort).toBe('low')
+    store.dropChatEffort('chat-1', latest)
     expect(useMothershipEffortStore.getState().chatEfforts).toEqual({})
   })
+
+  it('keeps an adopted new-chat pick when a stale save token fails', () => {
+    const store = useMothershipEffortStore.getState()
+    const stale = store.setChatEffort('chat-1', 'low')
+    store.adoptNewChatEffort('chat-1', 'low')
+    store.dropChatEffort('chat-1', stale)
+    expect(useMothershipEffortStore.getState().chatEfforts['chat-1']?.effort).toBe('low')
+  })
 })

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

@@ -7,6 +7,12 @@ import {
   resolveMothershipModelSettings,
 } from '@/lib/mothership/model-options'
 
+/** A chat's effort pick and the token that tells it apart from other picks of the same value. */
+interface ChatEffortPick {
+  effort: MothershipEffort
+  pick: number
+}
+
 interface MothershipEffortState {
   modelSelection: ModelSelection
   setModel: (model: ModelSelection['model']) => void
@@ -21,10 +27,11 @@ interface MothershipEffortState {
    * Picks made in existing chats this session, by chat id. They win over the chat's loaded
    * value, so a detail refetch or a save still in flight never shows or sends an older one.
    */
-  chatEfforts: Record<string, MothershipEffort>
-  setChatEffort: (chatId: string, effort: MothershipEffort) => void
-  /** Drops a pick whose save failed, unless a newer pick replaced it. */
-  dropChatEffort: (chatId: string, effort: MothershipEffort) => void
+  chatEfforts: Record<string, ChatEffortPick>
+  /** Records a pick and returns its token for {@link MothershipEffortState.dropChatEffort}. */
+  setChatEffort: (chatId: string, effort: MothershipEffort) => number
+  /** Drops a pick whose save failed, unless a newer pick replaced it, even one of the same value. */
+  dropChatEffort: (chatId: string, pick: number) => void
   /** Moves the new-chat pick onto the chat its first send created. */
   adoptNewChatEffort: (chatId: string, effort: MothershipEffort) => void
   reset: () => void
@@ -39,6 +46,9 @@ const initialState: Pick<
   chatEfforts: {},
 }
 
+/** Counts picks across chats so a token never repeats within a session. */
+let lastChatEffortPick = 0
+
 function withModelSelection(
   modelSelection: ModelSelection
 ): Pick<MothershipEffortState, 'modelSelection'> {
@@ -55,18 +65,23 @@ export const useMothershipEffortStore = create<MothershipEffortState>()(
         setModel: (model) =>
           set((state) => withModelSelection({ model, fastMode: state.modelSelection.fastMode })),
         setNewChatEffort: (newChatEffort) => set({ newChatEffort }),
-        setChatEffort: (chatId, effort) =>
-          set((state) => ({ chatEfforts: { ...state.chatEfforts, [chatId]: effort } })),
-        dropChatEffort: (chatId, effort) =>
+        setChatEffort: (chatId, effort) => {
+          const pick = ++lastChatEffortPick
+          set((state) => ({ chatEfforts: { ...state.chatEfforts, [chatId]: { effort, pick } } }))
+          return pick
+        },
+        dropChatEffort: (chatId, pick) =>
           set((state) => {
-            if (state.chatEfforts[chatId] !== effort) return state
+            if (state.chatEfforts[chatId]?.pick !== pick) return state
             return { chatEfforts: omit(state.chatEfforts, [chatId]) }
           }),
-        adoptNewChatEffort: (chatId, effort) =>
+        adoptNewChatEffort: (chatId, effort) => {
+          const pick = ++lastChatEffortPick
           set((state) => ({
             newChatEffort: null,
-            chatEfforts: { ...state.chatEfforts, [chatId]: effort },
-          })),
+            chatEfforts: { ...state.chatEfforts, [chatId]: { effort, pick } },
+          }))
+        },
         reset: () => set(initialState),
       }),
       {