fix(mothership): keep the latest effort pick through refetches, failed saves and abandoned new chats

Author: waleedlatif1

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

@@ -1,5 +1,6 @@
 'use client'
 
+import { useEffect } from 'react'
 import {
   DropdownMenu,
   DropdownMenuContent,
@@ -33,10 +34,13 @@ export function ModelSelector() {
   const setFastMode = useMothershipEffortStore((state) => state.setFastMode)
   const { chatId } = useChatSurface()
   const { data: chatHistory } = useMothershipChatHistory(chatId)
+  const chatPick = useMothershipEffortStore((state) =>
+    chatId ? state.chatEfforts[chatId] : undefined
+  )
   const newChatEffort = useMothershipEffortStore((state) => state.newChatEffort)
   const setNewChatEffort = useMothershipEffortStore((state) => state.setNewChatEffort)
-  const { mutate: setChatEffort } = useSetMothershipChatEffort()
-  const effortChoice = chatId ? chatHistory?.effort : newChatEffort
+  const { mutate: saveChatEffort } = useSetMothershipChatEffort(chatId)
+  const effortChoice = chatId ? (chatPick ?? chatHistory?.effort) : newChatEffort
   const { effort, modelSelection } = resolveMothershipModelSettings(
     { effort: effortChoice ?? undefined, modelSelection: selection },
     advanced
@@ -45,9 +49,15 @@ export function ModelSelector() {
     ? mothershipEffortOptions(modelSelection.model)
     : MOTHERSHIP_SIMPLE_EFFORT_OPTIONS
   const setEffort = (choice: MothershipEffort) => {
-    if (chatId) setChatEffort({ chatId, effort: choice })
+    if (chatId) saveChatEffort(choice)
     else setNewChatEffort(choice)
   }
+
+  useEffect(() => {
+    if (chatId) return
+    return () => useMothershipEffortStore.getState().setNewChatEffort(null)
+  }, [chatId])
+
   const effortLabel = options.find((option) => option.value === effort)?.label ?? effort
   const modelLabel =
     MOTHERSHIP_MODEL_OPTIONS.find((option) => option.value === modelSelection.model)?.label ??

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

@@ -3396,6 +3396,18 @@ export function useChat(
 
       let requestChatId =
         queuedSendHandoff?.chatId ?? selectedChatIdRef.current ?? chatIdRef.current
+      // Read before the composer can unmount. Sent only when picked; otherwise the server
+      // uses the chat's stored pick or the default.
+      const effortStore = useMothershipEffortStore.getState()
+      const effortChoice =
+        options?.requestMode === 'assistant'
+          ? undefined
+          : requestChatId
+            ? (effortStore.chatEfforts[requestChatId] ??
+              queryClient.getQueryData<MothershipChatHistory>(
+                mothershipChatKeys.detail(requestChatId)
+              )?.effort)
+            : effortStore.newChatEffort
       const writeQueuedSendHandoff = (chatId?: string) => {
         if (!queuedSendHandoff) return
         if (!chatId && !queuedSendHandoff.supersededStreamId) return
@@ -3681,12 +3693,6 @@ export function useChat(
             ? {}
             : await getDesktopChatCapabilities(desktopScopeIdRef.current)
 
-        // Sent only when picked; otherwise the server uses the chat's stored pick or the default.
-        const effortChoice = requestChatId
-          ? queryClient.getQueryData<MothershipChatHistory>(
-              mothershipChatKeys.detail(requestChatId)
-            )?.effort
-          : useMothershipEffortStore.getState().newChatEffort
         const response = await fetch(apiPathRef.current, {
           method: 'POST',
           headers: { 'Content-Type': 'application/json' },
@@ -3734,13 +3740,8 @@ export function useChat(
           return consumedByTranscript
         }
         if (admittedChatId && !requestChatId) {
-          if (effortChoice) {
-            queryClient.setQueryData<MothershipChatHistory>(
-              mothershipChatKeys.detail(admittedChatId),
-              (current) => current && { ...current, effort: effortChoice }
-            )
-            useMothershipEffortStore.getState().setNewChatEffort(null)
-          }
+          if (effortChoice)
+            useMothershipEffortStore.getState().adoptNewChatEffort(admittedChatId, effortChoice)
           requestChatId = admittedChatId
           streamTargetChatId = admittedChatId
           adoptResolvedChatId(admittedChatId, { replaceHomeHistory: true, invalidateList: true })
@@ -3789,6 +3790,8 @@ export function useChat(
             const conflictChatId =
               typeof errorData.chatId === 'string' ? errorData.chatId : undefined
             if (conflictChatId && !streamTargetChatId) {
+              if (effortChoice)
+                useMothershipEffortStore.getState().adoptNewChatEffort(conflictChatId, effortChoice)
               adoptResolvedChatId(conflictChatId, {
                 replaceHomeHistory: true,
                 invalidateList: true,

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

@@ -35,6 +35,7 @@ import {
 } from '@/lib/mothership/request/session/file-preview-session-contract'
 import { isStreamBatchEvent, type StreamBatchEvent } from '@/lib/mothership/request/session/types'
 import type { MothershipResource } from '@/lib/mothership/resources/types'
+import { useMothershipEffortStore } from '@/stores/mothership-effort/store'
 import { useMothershipQueueStore } from '@/stores/mothership-queue/store'
 
 export interface MothershipChatMetadata {
@@ -598,27 +599,27 @@ async function setChatEffort({
 }
 
 /**
- * Records the effort the user picked for a chat, shown immediately. Success needs no
- * refetch: the server stores exactly the picked value, and refetching the detail would
- * reload the whole transcript, possibly mid-stream.
+ * Records the effort the user picked for a chat. The pick shows and sends at once from the
+ * session's pick map; saves for one chat run one at a time so the last pick is the one stored.
  */
-export function useSetMothershipChatEffort() {
+export function useSetMothershipChatEffort(chatId: string | undefined) {
   const queryClient = useQueryClient()
   return useMutation({
-    mutationFn: setChatEffort,
-    onMutate: async ({ chatId, effort }) => {
-      const queryKey = mothershipChatKeys.detail(chatId)
-      const previousEffort = queryClient.getQueryData<MothershipChatHistory>(queryKey)?.effort
-      queryClient.setQueryData<MothershipChatHistory>(
-        queryKey,
-        (current) => current && { ...current, effort }
-      )
-      return { previousEffort }
+    mutationFn: (effort: MothershipEffort) => {
+      if (!chatId) throw new Error('A chat effort needs a chat')
+      return setChatEffort({ chatId, effort })
     },
-    onError: (_error, { chatId }, context) => {
+    scope: { id: `mothership-chat-effort:${chatId ?? ''}` },
+    onMutate: (effort) => {
+      if (chatId) useMothershipEffortStore.getState().setChatEffort(chatId, effort)
+    },
+    onError: (_error, effort) => {
+      if (chatId) useMothershipEffortStore.getState().dropChatEffort(chatId, effort)
+    },
+    onSuccess: (_data, effort) => {
       queryClient.setQueryData<MothershipChatHistory>(
         mothershipChatKeys.detail(chatId),
-        (current) => current && { ...current, effort: context?.previousEffort }
+        (current) => current && { ...current, effort }
       )
     },
   })

apps/sim/lib/mothership/chat/post.test.ts

@@ -845,6 +845,7 @@ describe('handleUnifiedChatPost', () => {
     { stored: 'high', sent: undefined, runs: 'high' },
     { stored: 'low', sent: undefined, runs: 'low' },
     { stored: 'high', sent: 'low', runs: 'low' },
+    { stored: 'high', sent: 'max', runs: 'xhigh' },
   ] as const)(
     'runs a chat whose stored effort choice is $stored at $runs when the send names $sent',
     async ({ stored, sent, runs }) => {
@@ -869,7 +870,7 @@ describe('handleUnifiedChatPost', () => {
       expect(buildCopilotRequestPayload).toHaveBeenCalledWith(
         expect.objectContaining({ effort: runs })
       )
-      expect(admitTurn.mock.calls[0][0].input.effortChoice).toBe(sent)
+      expect(admitTurn.mock.calls[0][0].input.effortChoice).toBe(sent && runs)
     }
   )
 

apps/sim/lib/mothership/chat/post.ts

@@ -1480,7 +1480,8 @@ export async function handleUnifiedChatPost(req: NextRequest) {
               requestMode: body.mode,
             },
             notifyWorkspaceStatus: branch.notifyChatStatus,
-            effortChoice,
+            // The effort this turn actually runs at, so the stored pick is always one it can use.
+            effortChoice: effortChoice && body.effort,
           },
         })
         // Admission committed. A failure to attach this HTTP sink must leave the turn recoverable.

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

@@ -39,4 +39,14 @@ describe('Build reasoning preferences', () => {
       modelSelection: { model: 'gpt-6-astra', fastMode: true },
     })
   })
+
+  it('keeps a newer chat pick when an older pick fails to save', () => {
+    const store = useMothershipEffortStore.getState()
+    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')
+    expect(useMothershipEffortStore.getState().chatEfforts).toEqual({})
+  })
 })

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

@@ -1,4 +1,4 @@
-import { toRecord } from '@sim/utils/object'
+import { omit, toRecord } from '@sim/utils/object'
 import { create } from 'zustand'
 import { devtools, persist } from 'zustand/middleware'
 import { type ModelSelection, ModelSelectionSchema } from '@/lib/mothership/generated/protocol'
@@ -13,17 +13,31 @@ interface MothershipEffortState {
   setFastMode: (fastMode: boolean) => void
   /**
    * The effort picked in a composer whose chat does not exist yet. Its first send records
-   * it on the new chat; existing chats keep their own choice on the chat itself.
+   * it on the new chat; leaving that composer unsent drops it.
    */
   newChatEffort: MothershipEffort | null
   setNewChatEffort: (effort: MothershipEffort | null) => void
+  /**
+   * 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
+  /** Moves the new-chat pick onto the chat its first send created. */
+  adoptNewChatEffort: (chatId: string, effort: MothershipEffort) => void
   reset: () => void
 }
 
-const initialState = {
+const initialState: Pick<
+  MothershipEffortState,
+  'modelSelection' | 'newChatEffort' | 'chatEfforts'
+> = {
   modelSelection: { model: 'gpt-6-astra', fastMode: false },
   newChatEffort: null,
-} satisfies Pick<MothershipEffortState, 'modelSelection' | 'newChatEffort'>
+  chatEfforts: {},
+}
 
 function withModelSelection(
   modelSelection: ModelSelection
@@ -41,6 +55,18 @@ 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) =>
+          set((state) => {
+            if (state.chatEfforts[chatId] !== effort) return state
+            return { chatEfforts: omit(state.chatEfforts, [chatId]) }
+          }),
+        adoptNewChatEffort: (chatId, effort) =>
+          set((state) => ({
+            newChatEffort: null,
+            chatEfforts: { ...state.chatEfforts, [chatId]: effort },
+          })),
         reset: () => set(initialState),
       }),
       {