fix(mothership): make chat fork complete, consistent and safe to retry

apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts

@@ -185,6 +185,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
     queueTableRows(copilotChats, [chat])
     queueTableRows(copilotChats, [chat])
     queueTableRows(member, [{ role: 'member' }])
+    queueTableRows(copilotChats, [{ id: chat.id }])
     const attachments = [
       {
         id: 'upload-1',
@@ -215,6 +216,26 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
     expect(mockAssertActiveWorkspaceAccess).not.toHaveBeenCalled()
   })
 
+  it('refuses to publish a fork whose source chat was purged while it copied', async () => {
+    // The fork shares its attachment keys with the source, and cleanup deletes an unreferenced
+    // key only after deleting the source row: publishing without the source would leave the
+    // fork pointing at bytes that cleanup is about to delete.
+    dbChainMockFns.limit.mockReset()
+    const chat = { ...parentRow, workspaceId: null, organizationId: 'org-1', resources: [] }
+    queueTableRows(copilotChats, [chat])
+    queueTableRows(copilotChats, [chat])
+    queueTableRows(member, [{ role: 'member' }])
+    queueTableRows(copilotChats, [])
+    const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(res.status).toBe(404)
+    expect(mockAppendCopilotChatMessages).not.toHaveBeenCalled()
+    expect(mockPublishStatusChanged).not.toHaveBeenCalled()
+    expect(mockFetchGo.mock.calls.map(([url]) => url)).toEqual([
+      'http://mothership.test/api/chats/fork',
+      'http://mothership.test/api/tasks/cleanup',
+    ])
+  })
+
   it.each(['membership', 'capability'] as const)(
     'denies organization forks after %s revocation before copying',
     async (revocation) => {
@@ -316,6 +337,12 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
     expect(mockPublishStatusChanged).not.toHaveBeenCalled()
   })
 
+  it.each([404, 409, 413])('passes a worker %i refusal through', async (status) => {
+    mockFetchGo.mockResolvedValue(Response.json({ error: 'refused' }, { status }))
+    const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(res.status).toBe(status)
+  })
+
   it('surfaces failed blob copies and excludes their metadata from publication', async () => {
     mockExecuteChatFileBlobCopies.mockResolvedValue({
       copied: 1,
@@ -336,6 +363,116 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
     expect(dbChainMockFns.delete).not.toHaveBeenCalled()
   })
 
+  it('discards the worker copy when the fork cannot be published', async () => {
+    dbChainMockFns.transaction.mockRejectedValueOnce(new Error('connection lost'))
+    const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(res.status).toBe(500)
+    const [fork, discard] = mockFetchGo.mock.calls
+    expect(fork[0]).toBe('http://mothership.test/api/chats/fork')
+    expect(discard[0]).toBe('http://mothership.test/api/tasks/cleanup')
+    expect(JSON.parse(discard[1].body)).toEqual({
+      chatIds: [JSON.parse(fork[1].body).newChatId],
+    })
+    expect(mockPublishStatusChanged).not.toHaveBeenCalled()
+  })
+
+  it('keeps references on the source file when its copy fails', async () => {
+    const oldKey = 'workspace/ws-1/old-cat.png'
+    mockListForkableChatFiles.mockResolvedValue([
+      { id: OLD_FILE_ID, key: oldKey, messageId: 'msg-1', workspaceId: 'ws-1' },
+    ])
+    mockPlanChatFileCopies.mockReturnValue({
+      idMap: new Map([[OLD_FILE_ID, NEW_FILE_ID]]),
+      keyMap: new Map([[oldKey, 'workspace/ws-1/new-cat.png']]),
+      blobTasks: [
+        {
+          copyId: NEW_FILE_ID,
+          sourceKey: oldKey,
+          targetKey: 'workspace/ws-1/new-cat.png',
+          context: 'mothership',
+          fileName: 'cat.png',
+          contentType: 'image/png',
+        },
+      ],
+    })
+    mockExecuteChatFileBlobCopies.mockResolvedValue({
+      copied: 0,
+      failed: 1,
+      failedCopyIds: [NEW_FILE_ID],
+    })
+    const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(res.status).toBe(200)
+    // The copy is never published, so neither Sim's transcript nor the worker's may name it.
+    expect(mockAppendCopilotChatMessages.mock.calls[0][1][0].content).toBe(
+      `See ![cat](/api/files/view/${OLD_FILE_ID})`
+    )
+    const forkRequest = JSON.parse(mockFetchGo.mock.calls[0][1].body)
+    expect(forkRequest.fileIds).toEqual({})
+    expect(forkRequest.fileKeys).toEqual({})
+  })
+
+  it('re-points in-app file links and tool-call arguments at the copied file', async () => {
+    mockLoadCopilotChatMessages.mockResolvedValue([
+      { ...threeMessages[0], content: `Open /workspace/ws-1/files/${OLD_FILE_ID}` },
+      {
+        ...threeMessages[1],
+        contentBlocks: [
+          {
+            type: 'tool',
+            toolCall: {
+              id: 'call-1',
+              name: 'sim_cli',
+              state: 'success',
+              params: { fileId: OLD_FILE_ID, args: ['files', 'read', OLD_FILE_ID], n: 2 },
+              display: { title: `Read /api/files/view/${OLD_FILE_ID}` },
+            },
+          },
+        ],
+      },
+    ])
+    mockPlanChatFileCopies.mockReturnValue({
+      idMap: new Map([[OLD_FILE_ID, NEW_FILE_ID]]),
+      keyMap: new Map(),
+      blobTasks: [],
+    })
+    const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(res.status).toBe(200)
+    const [user, assistant] = mockAppendCopilotChatMessages.mock.calls[0][1]
+    expect(user.content).toBe(`Open /workspace/ws-1/files/${NEW_FILE_ID}`)
+    expect(assistant.contentBlocks[0].toolCall).toMatchObject({
+      params: { fileId: NEW_FILE_ID, args: ['files', 'read', NEW_FILE_ID], n: 2 },
+      display: { title: `Read /api/files/view/${NEW_FILE_ID}` },
+    })
+  })
+
+  it('drops a Sources tab whose response is past the cut', async () => {
+    const kept = { type: 'sources', id: 'cited-sources', title: 'Sources' }
+    dbChainMockFns.limit.mockResolvedValue([
+      { ...parentRow, resources: [{ ...kept, sources: { messageId: 'msg-3' } }] },
+    ])
+    await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(dbChainMockFns.values).toHaveBeenCalledWith(expect.objectContaining({ resources: [] }))
+
+    dbChainMockFns.values.mockClear()
+    dbChainMockFns.limit.mockResolvedValue([
+      {
+        ...parentRow,
+        resources: [{ ...kept, sources: { messageId: 'live-id', requestId: 'req-2' } }],
+      },
+    ])
+    mockLoadCopilotChatMessages.mockResolvedValue([
+      threeMessages[0],
+      { ...threeMessages[1], requestId: 'req-2' },
+      threeMessages[2],
+    ])
+    await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
+    expect(dbChainMockFns.values).toHaveBeenCalledWith(
+      expect.objectContaining({
+        resources: [{ ...kept, sources: { messageId: 'live-id', requestId: 'req-2' } }],
+      })
+    )
+  })
+
   it('copies pre-cut uploads and drops only post-cut ghosts', async () => {
     // The source chat owns two more uploads (apple pre-cut, banana post-cut)
     // beside the kept one, plus one shared workspace-file resource. The fork

apps/sim/app/api/mothership/chats/[chatId]/fork/route.ts

@@ -2,7 +2,7 @@ import { createLogger } from '@sim/logger'
 import { type NextRequest, NextResponse } from 'next/server'
 import { forkMothershipChatContract } from '@/lib/api/contracts/mothership-chats'
 import { parseRequest } from '@/lib/api/server'
-import { asOrchestrationError } from '@/lib/core/orchestration/types'
+import { asOrchestrationError, statusForOrchestrationError } from '@/lib/core/orchestration/types'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { forkChat } from '@/lib/mothership/chat/application/fork'
 import {
@@ -38,6 +38,11 @@ export const POST = withRouteHandler(
       if (classified?.code === 'not_found' || classified?.code === 'forbidden')
         return NextResponse.json({ error: 'Chat not found' }, { status: 404 })
       if (classified?.code === 'validation') return createBadRequestResponse(classified.message)
+      if (classified?.code === 'conflict' || classified?.code === 'payload_too_large')
+        return NextResponse.json(
+          { error: classified.message },
+          { status: statusForOrchestrationError(classified.code) }
+        )
       logger.error('Error forking chat:', error)
       return createInternalServerErrorResponse('Failed to fork chat')
     }

apps/sim/app/workspace/[workspaceId]/components/message-actions/message-actions.tsx

@@ -19,6 +19,7 @@ import {
   useCopyToClipboard,
 } from '@sim/emcn'
 import { useParams, useRouter } from 'next/navigation'
+import { isApiClientError } from '@/lib/api/client/errors'
 import { isLiveAssistantMessageId } from '@/lib/mothership/chat/live-message-id'
 import { organizationRoutes } from '@/lib/navigation/paths'
 import { useChatSurface } from '@/app/workspace/[workspaceId]/home/components/chat-surface-context'
@@ -40,6 +41,9 @@ interface MessageActionsProps {
   messageId?: string
 }
 
+/** Fork refusals whose message tells the person what to do: the response is unfinished, or the chat is too long. */
+const FORK_REFUSAL_STATUSES = new Set([409, 413])
+
 export const MessageActions = memo(function MessageActions({
   content,
   getCopyContent,
@@ -143,8 +147,12 @@ export const MessageActions = memo(function MessageActions({
         useFolderStore.getState().clearChatSelection()
         router.push(`/workspace/${params.workspaceId}/chat/${result.id}`)
       }
-    } catch {
-      toast.error('Failed to fork chat')
+    } catch (error) {
+      toast.error(
+        isApiClientError(error) && FORK_REFUSAL_STATUSES.has(error.status)
+          ? error.message
+          : 'Failed to fork chat'
+      )
     }
   }
 

apps/sim/lib/cleanup/chat-cleanup.test.ts

@@ -0,0 +1,101 @@
+import { copilotChats, copilotMessages, workspaceFiles } from '@sim/db/schema'
+import { queueTableRows, resetDbChainMock } from '@sim/testing/mocks/database.mock'
+import { storageServiceMockFns } from '@sim/testing/mocks/storage-service.mock'
+import { uploadsMock, uploadsMockFns } from '@sim/testing/mocks/uploads.mock'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+vi.mock('@/lib/uploads', () => uploadsMock)
+
+import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup'
+import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-key'
+
+const chatId = '3f0c2a52-8a43-4d4b-9b5f-0b3c7a1e2d10'
+
+function attachment(key: string) {
+  return { id: 'wf_x', key, filename: 'x.png', media_type: 'image/png', size: 1 }
+}
+
+/**
+ * Purges one deleted chat whose message rows are `messages`, while `remaining` are the message
+ * rows other chats of the organization still hold; returns the deleted keys by context.
+ */
+async function purge(
+  messages: Record<string, unknown>[],
+  remaining: Record<string, unknown>[] = []
+) {
+  queueTableRows(workspaceFiles, [])
+  queueTableRows(
+    copilotMessages,
+    messages.map((content) => ({ chatId, content }))
+  )
+  const cleanup = await prepareChatCleanup([chatId], 'test')
+  queueTableRows(copilotChats, [])
+  queueTableRows(
+    copilotMessages,
+    remaining.map((content) => ({ content }))
+  )
+  await cleanup.execute()
+  return Object.fromEntries(
+    storageServiceMockFns.mockDeleteFiles.mock.calls.map(([keys, context]) => [context, keys])
+  )
+}
+
+describe('chat purge storage', () => {
+  beforeEach(() => {
+    resetDbChainMock()
+    uploadsMockFns.mockIsUsingCloudStorage.mockReturnValue(true)
+    storageServiceMockFns.mockDeleteFiles.mockReset()
+    storageServiceMockFns.mockDeleteFiles.mockResolvedValue({ deleted: 1, failed: [] })
+  })
+
+  it('never deletes a workspace attachment as copilot storage', async () => {
+    // A fork carries its source's key when the file's copy failed or the file was deleted, and
+    // the copilot bucket can be the workspace bucket, so a workspace key here is another chat's file.
+    const deleted = await purge([
+      {
+        role: 'user',
+        content: 'look',
+        fileAttachments: [
+          attachment('workspace/ws-1/1700-abc-shared.png'),
+          attachment('copilot/1234/legacy.png'),
+        ],
+      },
+    ])
+    expect(deleted).toEqual({ copilot: ['copilot/1234/legacy.png'] })
+  })
+
+  it('deletes an organization attachment only once no remaining chat references it', async () => {
+    const shared = 'assistant/org-1/user-1/u1/shared.png'
+    const own = 'assistant/org-1/user-1/u2/own.png'
+    const message = {
+      role: 'user',
+      content: 'look',
+      fileAttachments: [attachment(shared), attachment(own)],
+    }
+    // A fork of this chat still holds the shared upload.
+    const deleted = await purge(
+      [message],
+      [{ role: 'user', content: 'fork', fileAttachments: [attachment(shared)] }]
+    )
+    expect(deleted).toEqual({ mothership: [own] })
+  })
+
+  it('deletes the inline images an assistant message published', async () => {
+    const deleted = await purge([
+      {
+        role: 'assistant',
+        requestId: 'req-1',
+        content: 'Here: ![chart](files/chart.png) and `![code](files/no.png)`',
+        contentBlocks: [{ type: 'text', content: '![other](/tmp/out.png)' }],
+      },
+      { role: 'user', requestId: 'req-2', content: '![u](files/u.png)' },
+      { role: 'assistant', content: '![x](files/x.png)' },
+    ])
+    expect(deleted).toEqual({
+      mothership: [
+        inlineChatImageKey(chatId, 'req-1', 'files/chart.png'),
+        inlineChatImageKey(chatId, 'req-1', '/tmp/out.png'),
+      ],
+    })
+  })
+})