fix(mothership): persist @-mentioned resources across send and merge on hydration

Author: waleedlatif1

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

@@ -5,6 +5,13 @@ import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { useQueryClient } from '@tanstack/react-query'
 import { usePathname, useRouter } from 'next/navigation'
+import { requestJson } from '@/lib/api/client/request'
+import {
+  addMothershipChatResourceContract,
+  removeMothershipChatResourceContract,
+  reorderMothershipChatResourcesContract,
+} from '@/lib/api/contracts/mothership-tasks'
+import { cancelWorkflowExecutionContract } from '@/lib/api/contracts/workflows'
 import { getMothershipAttachmentPreviewUrl } from '@/lib/copilot/chat/attachment-preview'
 import { toDisplayMessage } from '@/lib/copilot/chat/display-message'
 import { getLiveAssistantMessageId } from '@/lib/copilot/chat/effective-transcript'
@@ -1536,6 +1543,7 @@ export function useChat(
   }, [])
   const resourcesRef = useRef(resources)
   resourcesRef.current = resources
+  const pendingPersistResourceKeysRef = useRef<Set<string>>(new Set())
 
   // Derive the effective active resource ID — auto-selects the last resource when the stored ID is
   // absent or no longer in the list, avoiding a separate Effect-based state correction loop.
@@ -1962,6 +1970,7 @@ export function useChat(
     setTransportIdle()
     setResources([])
     setActiveResourceId(null)
+    pendingPersistResourceKeysRef.current.clear()
     resetEphemeralPreviewState()
     setMessageQueue([])
     clearQueueDispatchState()
@@ -1974,6 +1983,22 @@ export function useChat(
     setTransportIdle,
   ])
 
+  const flushPendingResources = useCallback((chatId: string) => {
+    const pendingKeys = pendingPersistResourceKeysRef.current
+    if (pendingKeys.size === 0) return
+    for (const resource of resourcesRef.current) {
+      const key = `${resource.type}:${resource.id}`
+      if (!pendingKeys.has(key)) continue
+      pendingKeys.delete(key)
+      requestJson(addMothershipChatResourceContract, { body: { chatId, resource } }).catch(
+        (err) => {
+          pendingPersistResourceKeysRef.current.add(key)
+          logger.warn('Failed to flush pending resource; will retry on next hydration', err)
+        }
+      )
+    }
+  }, [])
+
   const adoptResolvedChatId = useCallback(
     (chatId: string, options?: { replaceHomeHistory?: boolean; invalidateList?: boolean }) => {
       const selectedChatId = selectedChatIdRef.current
@@ -1992,8 +2017,9 @@ export function useChat(
       if (options?.invalidateList) {
         queryClient.invalidateQueries({ queryKey: taskKeys.list(workspaceId) })
       }
+      flushPendingResources(chatId)
     },
-    [queryClient, workspaceId]
+    [flushPendingResources, queryClient, workspaceId]
   )
 
   const { data: chatHistory } = useChatHistory(resolvedChatId)
@@ -2018,15 +2044,16 @@ export function useChat(
     }
 
     const persistChatId = chatIdRef.current ?? selectedChatIdRef.current
+    const key = `${resource.type}:${resource.id}`
     if (persistChatId) {
-      // boundary-raw-fetch: fire-and-forget side-effect during stream lifecycle; intentionally avoids requestJson's response parsing/throw semantics so a failure here cannot interrupt the active turn
-      fetch('/api/mothership/chat/resources', {
-        method: 'POST',
-        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({ chatId: persistChatId, resource }),
+      requestJson(addMothershipChatResourceContract, {
+        body: { chatId: persistChatId, resource },
       }).catch((err) => {
-        logger.warn('Failed to persist resource', err)
+        pendingPersistResourceKeysRef.current.add(key)
+        logger.warn('Failed to persist resource; will retry on next hydration', err)
       })
+    } else {
+      pendingPersistResourceKeysRef.current.add(key)
     }
     return true
   }, [])
@@ -2035,13 +2062,14 @@ export function useChat(
     setResources((prev) => prev.filter((r) => !(r.type === resourceType && r.id === resourceId)))
     setActiveResourceId((prev) => (prev === resourceId ? null : prev))
 
+    const key = `${resourceType}:${resourceId}`
+    const wasPending = pendingPersistResourceKeysRef.current.delete(key)
+    if (wasPending) return
+
     const persistChatId = chatIdRef.current ?? selectedChatIdRef.current
     if (persistChatId) {
-      // boundary-raw-fetch: fire-and-forget side-effect; intentionally avoids requestJson's response parsing/throw semantics so a transient failure cannot interrupt the caller
-      fetch('/api/mothership/chat/resources', {
-        method: 'DELETE',
-        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({ chatId: persistChatId, resourceType, resourceId }),
+      requestJson(removeMothershipChatResourceContract, {
+        body: { chatId: persistChatId, resourceType, resourceId },
       }).catch((err) => {
         logger.warn('Failed to persist resource removal', err)
       })
@@ -2050,6 +2078,18 @@ export function useChat(
 
   const reorderResources = useCallback((newOrder: MothershipResource[]) => {
     setResources(newOrder)
+    const persistChatId = chatIdRef.current ?? selectedChatIdRef.current
+    if (!persistChatId) return
+    const pendingKeys = pendingPersistResourceKeysRef.current
+    const persistableResources = newOrder.filter(
+      (r) => r.id !== 'streaming-file' && !pendingKeys.has(`${r.type}:${r.id}`)
+    )
+    if (persistableResources.length === 0) return
+    requestJson(reorderMothershipChatResourcesContract, {
+      body: { chatId: persistChatId, resources: persistableResources },
+    }).catch((err) => {
+      logger.warn('Failed to persist resource reorder', err)
+    })
   }, [])
 
   const ensureWorkflowToolResource = useCallback(
@@ -2179,6 +2219,7 @@ export function useChat(
     setTransportIdle()
     setResources([])
     setActiveResourceId(null)
+    pendingPersistResourceKeysRef.current.clear()
     resetEphemeralPreviewState()
     setMessageQueue([])
     clearQueueDispatchState()
@@ -2229,27 +2270,32 @@ export function useChat(
 
     const hasPersistedStreamingFile = chatHistory.resources.some((r) => r.id === 'streaming-file')
     if (hasPersistedStreamingFile) {
-      // boundary-raw-fetch: fire-and-forget cleanup during chat-history hydration; failures are silently swallowed to keep hydration non-blocking
-      fetch('/api/mothership/chat/resources', {
-        method: 'DELETE',
-        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({
+      requestJson(removeMothershipChatResourceContract, {
+        body: {
           chatId: chatHistory.id,
           resourceType: 'file',
           resourceId: 'streaming-file',
-        }),
+        },
       }).catch(() => {})
     }
 
+    flushPendingResources(chatHistory.id)
+
     const persistedResources = chatHistory.resources.filter((r) => r.id !== 'streaming-file')
-    if (persistedResources.length > 0) {
+    const serverKeys = new Set(persistedResources.map((r) => `${r.type}:${r.id}`))
+    const localOnly = resourcesRef.current.filter(
+      (r) => r.id !== 'streaming-file' && !serverKeys.has(`${r.type}:${r.id}`)
+    )
+    const mergedResources = [...persistedResources, ...localOnly]
+
+    if (mergedResources.length > 0) {
       const hydratedActiveResourceId =
         activeResourceIdRef.current &&
-        persistedResources.some((resource) => resource.id === activeResourceIdRef.current)
+        mergedResources.some((resource) => resource.id === activeResourceIdRef.current)
           ? activeResourceIdRef.current
-          : persistedResources[persistedResources.length - 1].id
+          : mergedResources[mergedResources.length - 1].id
       activeResourceIdRef.current = hydratedActiveResourceId
-      setResources(persistedResources)
+      setResources(mergedResources)
       setActiveResourceId(hydratedActiveResourceId)
 
       for (const resource of persistedResources) {
@@ -2373,6 +2419,7 @@ export function useChat(
     workspaceId,
     cancelActiveStreamReader,
     cancelActiveStreamRecovery,
+    flushPendingResources,
     queryClient,
     recoverPendingClientWorkflowTools,
     seedPreviewSessions,
@@ -5003,9 +5050,8 @@ export function useChat(
       const executionId = execState.getCurrentExecutionId(workflowId)
       if (executionId) {
         execState.setCurrentExecutionId(workflowId, null)
-        // boundary-raw-fetch: fire-and-forget execution cancellation invoked from a stop-generation barrier; failures are silently swallowed so the stop teardown cannot stall on a contract-validation throw
-        fetch(`/api/workflows/${workflowId}/executions/${executionId}/cancel`, {
-          method: 'POST',
+        requestJson(cancelWorkflowExecutionContract, {
+          params: { id: workflowId, executionId },
         }).catch(() => {})
       }