fix: harden batch add block enrichment fallback

Author: PlaneInABottle

apps/sim/socket/database/operations.test.ts

@@ -2,18 +2,39 @@
  * @vitest-environment node
  */
 
-import { describe, expect, it, vi } from 'vitest'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
 import {
   filterBatchAddBlocksByProtection,
   getBatchUpdateParentCoMovingKey,
+  persistWorkflowOperation,
 } from '@/socket/database/operations'
 
+const {
+  mockTransaction,
+  mockGetActiveWorkflowContext,
+  mockMergeSubBlockValues,
+  mockGetBlock,
+  mockWorkflowTable,
+  mockWorkflowBlocksTable,
+  mockWorkflowEdgesTable,
+  mockWorkflowSubflowsTable,
+} = vi.hoisted(() => ({
+  mockTransaction: vi.fn(),
+  mockGetActiveWorkflowContext: vi.fn(),
+  mockMergeSubBlockValues: vi.fn(),
+  mockGetBlock: vi.fn(),
+  mockWorkflowTable: {},
+  mockWorkflowBlocksTable: {},
+  mockWorkflowEdgesTable: {},
+  mockWorkflowSubflowsTable: {},
+}))
+
 vi.mock('@sim/db', () => ({
   webhook: {},
-  workflow: {},
-  workflowBlocks: {},
-  workflowEdges: {},
-  workflowSubflows: {},
+  workflow: mockWorkflowTable,
+  workflowBlocks: mockWorkflowBlocksTable,
+  workflowEdges: mockWorkflowEdgesTable,
+  workflowSubflows: mockWorkflowSubflowsTable,
 }))
 
 vi.mock('@sim/logger', () => ({
@@ -35,7 +56,7 @@ vi.mock('drizzle-orm', () => ({
 }))
 
 vi.mock('drizzle-orm/postgres-js', () => ({
-  drizzle: vi.fn(() => ({ transaction: vi.fn() })),
+  drizzle: vi.fn(() => ({ transaction: mockTransaction })),
 }))
 
 vi.mock('postgres', () => ({
@@ -59,26 +80,36 @@ vi.mock('@/lib/webhooks/provider-subscriptions', () => ({
 }))
 
 vi.mock('@/lib/workflows/active-context', () => ({
-  getActiveWorkflowContext: vi.fn(),
+  getActiveWorkflowContext: mockGetActiveWorkflowContext,
 }))
 
 vi.mock('@/lib/workflows/persistence/utils', () => ({
   loadWorkflowFromNormalizedTables: vi.fn(),
 }))
 
 vi.mock('@/lib/workflows/subblocks', () => ({
-  mergeSubBlockValues: vi.fn(),
+  mergeSubBlockValues: mockMergeSubBlockValues,
 }))
 
 vi.mock('@/blocks/registry', () => ({
-  getBlock: vi.fn(),
+  getBlock: mockGetBlock,
 }))
 
 vi.mock('@/triggers', () => ({
   getTrigger: vi.fn(),
   isTriggerValid: vi.fn(() => false),
 }))
 
+beforeEach(() => {
+  vi.clearAllMocks()
+  mockGetActiveWorkflowContext.mockResolvedValue({ id: 'workflow-1' })
+  mockMergeSubBlockValues.mockImplementation((subBlocks, values) => ({
+    ...((subBlocks as Record<string, unknown>) ?? {}),
+    ...((values as Record<string, unknown>) ?? {}),
+  }))
+  mockGetBlock.mockReturnValue(undefined)
+})
+
 describe('getBatchUpdateParentCoMovingKey', () => {
   it('groups blocks entering the same destination container together', () => {
     expect(getBatchUpdateParentCoMovingKey('old-parent-a', 'loop-1')).toBe(
@@ -161,3 +192,106 @@ describe('filterBatchAddBlocksByProtection', () => {
     expect(allowedBlocks.map((block) => block.id)).toEqual(['new-container', 'new-child'])
   })
 })
+
+describe('persistWorkflowOperation', () => {
+  it('defensively enriches batch-added blocks before persistence and broadcast payloads', async () => {
+    const insertedBlockValues: Array<Record<string, unknown>> = []
+
+    const tx = {
+      update: vi.fn(() => ({
+        set: vi.fn(() => ({ where: vi.fn().mockResolvedValue(undefined) })),
+      })),
+      select: vi.fn(() => ({
+        from: vi.fn(() => ({
+          where: vi.fn().mockResolvedValue([]),
+        })),
+      })),
+      insert: vi.fn((table) => {
+        if (table === mockWorkflowBlocksTable) {
+          return {
+            values: vi.fn((values) => {
+              insertedBlockValues.push(...(values as Array<Record<string, unknown>>))
+              return {
+                onConflictDoUpdate: vi.fn().mockResolvedValue(undefined),
+              }
+            }),
+          }
+        }
+
+        return {
+          values: vi.fn(() => ({
+            onConflictDoUpdate: vi.fn().mockResolvedValue(undefined),
+          })),
+        }
+      }),
+    }
+
+    mockTransaction.mockImplementation(async (callback) => callback(tx))
+    mockGetBlock.mockReturnValue({
+      subBlocks: [
+        { id: 'prompt', type: 'short-input', defaultValue: 'registry-default' },
+        { id: 'temperature', type: 'slider', defaultValue: 0.7 },
+      ],
+      outputs: {
+        result: { type: 'string' },
+      },
+    })
+
+    const result = await persistWorkflowOperation('workflow-1', {
+      operation: 'batch-add-blocks',
+      target: 'blocks',
+      payload: {
+        blocks: [
+          {
+            id: 'block-1',
+            type: 'agent',
+            name: 'Agent',
+            position: { x: 10, y: 20 },
+            data: {},
+          },
+        ],
+        edges: [],
+        loops: {},
+        parallels: {},
+        subBlockValues: {
+          'block-1': {
+            prompt: { value: 'user value' },
+          },
+        },
+      },
+      timestamp: Date.now(),
+      userId: 'user-1',
+    })
+
+    expect(insertedBlockValues).toHaveLength(1)
+    expect(insertedBlockValues[0]).toMatchObject({
+      id: 'block-1',
+      subBlocks: {
+        prompt: { value: 'user value' },
+        temperature: { id: 'temperature', type: 'slider', value: 0.7 },
+      },
+      outputs: {
+        result: { type: 'string' },
+      },
+    })
+    expect(result.appliedPayload).toMatchObject({
+      blocks: [
+        {
+          id: 'block-1',
+          subBlocks: {
+            prompt: { value: 'user value' },
+            temperature: { id: 'temperature', type: 'slider', value: 0.7 },
+          },
+          outputs: {
+            result: { type: 'string' },
+          },
+        },
+      ],
+      subBlockValues: {
+        'block-1': {
+          prompt: { value: 'user value' },
+        },
+      },
+    })
+  })
+})

apps/sim/socket/database/operations.ts

@@ -1130,6 +1130,9 @@ async function handleBlocksOperationTx(
         parallelCount: Object.keys(parallels || {}).length,
       })
 
+      let allowedSubBlockValues: Record<string, unknown> = {}
+      let enrichedAllowedBlocks: Array<Record<string, unknown>> = []
+
       if (blocks && blocks.length > 0) {
         // Fetch existing blocks to check for locked parents
         const existingBlocks = await tx
@@ -1167,9 +1170,23 @@ async function handleBlocksOperationTx(
           }
         }
 
-        // Blocks are pre-enriched by enrichBatchAddBlocksPayload() upstream,
-        // so subBlocks, outputs, and triggerMode are already populated.
-        const blockValues = allowedBlocks.map((block: Record<string, unknown>) => {
+        const allowedBlockIds = new Set(allowedBlocks.map((block) => block.id as string))
+        allowedSubBlockValues = Object.fromEntries(
+          Object.entries((subBlockValues as Record<string, unknown> | undefined) ?? {}).filter(
+            ([id]) => allowedBlockIds.has(id)
+          )
+        )
+
+        const enrichedAllowedPayload = enrichBatchAddBlocksPayload({
+          ...payload,
+          blocks: allowedBlocks,
+          subBlockValues: allowedSubBlockValues,
+        })
+
+        enrichedAllowedBlocks =
+          (enrichedAllowedPayload.blocks as Array<Record<string, unknown>> | undefined) ?? []
+
+        const blockValues = enrichedAllowedBlocks.map((block: Record<string, unknown>) => {
           return {
             id: block.id as string,
             workflowId,
@@ -1215,7 +1232,7 @@ async function handleBlocksOperationTx(
         // Create subflow entries for loop/parallel blocks (skip if already in payload)
         const loopIds = new Set(loops ? Object.keys(loops) : [])
         const parallelIds = new Set(parallels ? Object.keys(parallels) : [])
-        for (const block of allowedBlocks) {
+        for (const block of enrichedAllowedBlocks) {
           const blockId = block.id as string
           if (block.type === 'loop' && !loopIds.has(blockId)) {
             await tx
@@ -1262,7 +1279,7 @@ async function handleBlocksOperationTx(
 
         // Update parent subflow node lists
         const parentIds = new Set<string>()
-        for (const block of allowedBlocks) {
+        for (const block of enrichedAllowedBlocks) {
           const parentId = (block.data as Record<string, unknown>)?.parentId as string | undefined
           if (parentId) {
             parentIds.add(parentId)
@@ -1273,7 +1290,7 @@ async function handleBlocksOperationTx(
         }
       }
 
-      const allowedBlockIds = new Set(allowedBlocks.map((block) => block.id as string))
+      const allowedBlockIds = new Set(enrichedAllowedBlocks.map((block) => block.id as string))
       const validEdgeEndpointIds = new Set([...existingBlockIds, ...allowedBlockIds])
       const allowedEdges = ((edges as Array<Record<string, unknown>> | undefined) ?? []).filter(
         (edge) => {
@@ -1297,11 +1314,6 @@ async function handleBlocksOperationTx(
           allowedBlockIds.has(id)
         )
       )
-      const allowedSubBlockValues = Object.fromEntries(
-        Object.entries((subBlockValues as Record<string, unknown> | undefined) ?? {}).filter(
-          ([id]) => allowedBlockIds.has(id)
-        )
-      )
 
       if (allowedEdges.length > 0) {
         const edgeValues = allowedEdges.map((edge: Record<string, unknown>) => ({
@@ -1372,7 +1384,7 @@ async function handleBlocksOperationTx(
       return {
         appliedPayload: {
           ...payload,
-          blocks: allowedBlocks,
+          blocks: enrichedAllowedBlocks,
           edges: allowedEdges,
           loops: allowedLoops,
           parallels: allowedParallels,