fix(mothership): keep a new chat's effort pick changed while its first send is pending (#8662)

* fix(mothership): keep a new chat's effort pick changed while its first send is pending

* fix(mothership): keep effort saves per chat and save the pick on dedup adoption

* fix(mothership): detach the effort save observer once it settles

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx

@@ -2000,6 +2000,59 @@ describe('useChat remount send recovery', () => {
     expect(state.postBodies).toHaveLength(0)
   })
 
+  it.each([
+    { pendingPick: 'high', kept: 'high', saved: true },
+    { pendingPick: 'low', kept: 'low', saved: false },
+  ] as const)(
+    'keeps a new chat effort picked while its first send is pending ($pendingPick)',
+    async ({ pendingPick, kept, saved }) => {
+      mockRequestJson.mockClear()
+      useMothershipEffortStore.getState().reset()
+      useMothershipEffortStore.getState().setNewChatEffort('low')
+      const { getResult } = renderUseChat()
+      await act(async () => {
+        void getResult().sendMessage('Pick while pending')
+      })
+      await waitFor(() => state.postBodies.length === 1)
+      expect(state.postBodies[0]).toMatchObject({ effort: 'low' })
+
+      useMothershipEffortStore.getState().setNewChatEffort(pendingPick)
+      const userMessageId = state.postBodies[0].userMessageId ?? ''
+      await act(async () => {
+        state.pendingAdmissions.get(userMessageId)?.()
+      })
+      await waitFor(() => !getResult().isSending)
+
+      expect(useMothershipEffortStore.getState().chatEfforts[DEDUPED_CHAT_ID]?.effort).toBe(kept)
+      const saves = mockRequestJson.mock.calls.filter(
+        ([contract]) => contract.path === '/api/mothership/chats/[chatId]/effort'
+      )
+      expect(saves.map(([, input]) => input)).toEqual(
+        saved ? [{ params: { chatId: DEDUPED_CHAT_ID }, body: { effort: kept } }] : []
+      )
+    }
+  )
+
+  it('saves the latest new-chat effort to the chat a deduplicated send names', async () => {
+    mockRequestJson.mockClear()
+    useMothershipEffortStore.getState().reset()
+    useMothershipEffortStore.getState().setNewChatEffort('high')
+    state.postBehavior = 'deduped'
+    const { getResult } = renderUseChat()
+    await act(async () => {
+      void getResult().sendMessage('Retry of an admitted send')
+    })
+    await waitFor(() => state.postBodies.length === 1 && !getResult().isSending)
+
+    expect(useMothershipEffortStore.getState().chatEfforts[DEDUPED_CHAT_ID]?.effort).toBe('high')
+    const saves = mockRequestJson.mock.calls.filter(
+      ([contract]) => contract.path === '/api/mothership/chats/[chatId]/effort'
+    )
+    expect(saves.map(([, input]) => input)).toEqual([
+      { params: { chatId: DEDUPED_CHAT_ID }, body: { effort: 'high' } },
+    ])
+  })
+
   it('loads the saved transcript once when its own stream completes', async () => {
     const chatId = 'chat-own-completion'
     const history: MothershipChatHistory = {

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

@@ -115,6 +115,7 @@ import {
   fetchMothershipChatHistory,
   type MothershipChatHistory,
   mothershipChatKeys,
+  saveMothershipChatEffort,
   useMothershipChatHistory,
 } from '@/hooks/queries/mothership-chats'
 import { fetchWorkflowEnvelope } from '@/hooks/queries/utils/fetch-workflow-envelope'
@@ -3435,6 +3436,18 @@ export function useChat(
                 mothershipChatKeys.detail(requestChatId)
               )?.effort)
             : effortStore.newChatEffort
+      /* Moves the new-chat pick onto the chat a send opened and saves it there unless it is
+         the pick this send's admission stored. A pick changed while the send was pending,
+         or a chat an earlier attempt opened with an unknown pick, gets the latest one. */
+      const adoptNewChatEffort = (chatId: string, admittedThisSend: boolean) => {
+        if (options?.requestMode === 'assistant') return
+        const store = useMothershipEffortStore.getState()
+        const latestChoice = store.newChatEffort ?? effortChoice
+        if (!latestChoice) return
+        store.adoptNewChatEffort(chatId, latestChoice)
+        if (!admittedThisSend || latestChoice !== effortChoice)
+          saveMothershipChatEffort(queryClient, chatId, latestChoice)
+      }
       const writeQueuedSendHandoff = (chatId?: string) => {
         if (!queuedSendHandoff) return
         if (!chatId && !queuedSendHandoff.supersededStreamId) return
@@ -3777,8 +3790,7 @@ export function useChat(
           return consumedByTranscript
         }
         if (admittedChatId && !requestChatId) {
-          if (effortChoice)
-            useMothershipEffortStore.getState().adoptNewChatEffort(admittedChatId, effortChoice)
+          adoptNewChatEffort(admittedChatId, true)
           requestChatId = admittedChatId
           streamTargetChatId = admittedChatId
           adoptResolvedChatId(admittedChatId, { replaceHomeHistory: true, invalidateList: true })
@@ -3827,9 +3839,7 @@ export function useChat(
             const conflictChatId =
               typeof errorData.chatId === 'string' ? errorData.chatId : undefined
             if (conflictChatId && !streamTargetChatId) {
-              // The retry carries the same pick the first attempt stored on that chat.
-              if (effortChoice)
-                useMothershipEffortStore.getState().adoptNewChatEffort(conflictChatId, effortChoice)
+              adoptNewChatEffort(conflictChatId, false)
               adoptResolvedChatId(conflictChatId, {
                 replaceHomeHistory: true,
                 invalidateList: true,

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

@@ -2,6 +2,9 @@ import { toError } from '@sim/utils/errors'
 import { isRecordLike } from '@sim/utils/object'
 import {
   keepPreviousData,
+  MutationObserver,
+  mutationOptions,
+  type QueryClient,
   queryOptions,
   skipToken,
   useMutation,
@@ -604,7 +607,28 @@ async function setChatEffort({
  */
 export function useSetMothershipChatEffort(chatId: string | undefined) {
   const queryClient = useQueryClient()
-  return useMutation({
+  return useMutation(chatEffortMutationOptions(queryClient, chatId))
+}
+
+/**
+ * Saves a pick for a chat learned outside render, such as the chat a send just opened. It
+ * shares the hook's per-chat scope, so it lands in order with picks made in the composer.
+ */
+export function saveMothershipChatEffort(
+  queryClient: QueryClient,
+  chatId: string,
+  effort: MothershipEffort
+): void {
+  const observer = new MutationObserver(queryClient, chatEffortMutationOptions(queryClient, chatId))
+  // Detaching once the save settles lets the mutation cache collect it.
+  observer
+    .mutate(effort)
+    .catch(() => undefined)
+    .finally(() => observer.reset())
+}
+
+function chatEffortMutationOptions(queryClient: QueryClient, chatId: string | undefined) {
+  return mutationOptions({
     mutationFn: (effort: MothershipEffort) => {
       if (!chatId) throw new Error('A chat effort needs a chat')
       return setChatEffort({ chatId, effort })