acquire MCP server locks before reads in attachment reconcile (TOCTOU)

Author: icecrasher321

apps/sim/lib/workspaces/fork/copy/workflow-mcp-attachments.test.ts

@@ -26,10 +26,11 @@ function makeTx(selectResults: unknown[][]) {
   const inserted: Array<Record<string, unknown>> = []
   const updates: Array<Record<string, unknown>> = []
   let call = 0
+  const select = vi.fn(() => ({
+    from: () => ({ where: () => Promise.resolve(selectResults[call++] ?? []) }),
+  }))
   const tx = {
-    select: () => ({
-      from: () => ({ where: () => Promise.resolve(selectResults[call++] ?? []) }),
-    }),
+    select,
     insert: () => ({
       values: (values: Array<Record<string, unknown>>) => {
         inserted.push(...values)
@@ -45,7 +46,7 @@ function makeTx(selectResults: unknown[][]) {
       }),
     }),
   }
-  return { tx: tx as unknown as DbOrTx, inserted, updates }
+  return { tx: tx as unknown as DbOrTx, inserted, updates, select }
 }
 
 const attachment = (overrides: Record<string, unknown> = {}) => ({
@@ -69,7 +70,8 @@ const serverMappingRow = {
 describe('reconcileForkWorkflowMcpAttachments', () => {
   it('creates the target attachment for a mapped server + written pair (push: child -> parent)', async () => {
     mockGetEdgeMappingRows.mockResolvedValue([serverMappingRow])
-    const { tx, inserted } = makeTx([
+    mockAcquireLock.mockClear()
+    const { tx, inserted, select } = makeTx([
       [{ id: 'srv-child' }, { id: 'srv-parent' }], // both mapped servers still live
       [attachment({ serverId: 'srv-child', workflowId: 'wf-child' })],
       [], // no existing target attachments
@@ -89,6 +91,11 @@ describe('reconcileForkWorkflowMcpAttachments', () => {
     })
     expect(result.affectedServerIds).toEqual(['srv-parent'])
     expect(mockAcquireLock).toHaveBeenCalledWith(tx, 'srv-parent')
+    // The lock must precede every read: locking after the diff is computed would let a
+    // concurrent attach commit in between and abort the promote on the unique constraint.
+    expect(mockAcquireLock.mock.invocationCallOrder[0]).toBeLessThan(
+      select.mock.invocationCallOrder[0]
+    )
   })
 
   it('archives a target attachment whose source counterpart was detached', async () => {

apps/sim/lib/workspaces/fork/copy/workflow-mcp-attachments.ts

@@ -122,10 +122,21 @@ export async function reconcileForkWorkflowMcpAttachments(params: {
   }
   if (serverMap.size === 0) return { affectedServerIds: [] }
 
+  // Same per-server serialization as the deploy-time tool sync and the attach/delete routes,
+  // in sorted order so two concurrent syncs can't deadlock on each other's server locks.
+  // Acquired BEFORE the reads below: every other attachment writer locks first, so locking
+  // after computing the diff would let a concurrent attach commit in between and turn our
+  // insert into a unique-constraint abort of the whole promote transaction (and a concurrent
+  // server delete into an FK abort).
+  for (const serverId of [...new Set(serverMap.values())].sort()) {
+    await acquireWorkflowMcpServerLock(tx, serverId)
+  }
+
   // Liveness guard: a mapped server may have been deleted since the fork (server deletion is a
   // hard delete that cascades its tools but leaves the identity row). A dead SOURCE server has
   // nothing to mirror; a dead TARGET server must be skipped or the insert below would violate
-  // the `server_id` FK and abort the whole promote transaction.
+  // the `server_id` FK and abort the whole promote transaction. Target liveness is stable for
+  // the rest of the transaction: the delete route takes the per-server lock we now hold.
   const mappedServerIds = [...new Set([...serverMap.keys(), ...serverMap.values()])]
   const liveServerIds = new Set(
     (
@@ -270,12 +281,6 @@ export async function reconcileForkWorkflowMcpAttachments(params: {
     return { affectedServerIds: [] }
   }
 
-  // Same per-server serialization as the deploy-time tool sync, in sorted order so two
-  // concurrent syncs can't deadlock on each other's server locks.
-  for (const serverId of [...affectedServerIds].sort()) {
-    await acquireWorkflowMcpServerLock(tx, serverId)
-  }
-
   if (inserts.length > 0) await tx.insert(workflowMcpTool).values(inserts)
   for (const update of updates) {
     await tx.update(workflowMcpTool).set(update.set).where(eq(workflowMcpTool.id, update.id))