cleanup migration code

icecrasher321 · icecrasher321 · commit 6489916ad868 · 2026-03-05T12:16:14.000-08:00
diff --git a/apps/sim/lib/workflows/migrations/subblock-migrations.test.ts b/apps/sim/lib/workflows/migrations/subblock-migrations.test.ts
@@ -20,7 +20,7 @@ function makeBlock(overrides: Partial<BlockState> & { type: string }): BlockStat
 describe('migrateSubblockIds', () => {
   describe('knowledge block', () => {
     it('should rename knowledgeBaseId to knowledgeBaseSelector', () => {
-      const blocks: Record<string, BlockState> = {
+      const input: Record<string, BlockState> = {
         b1: makeBlock({
           type: 'knowledge',
           subBlocks: {
@@ -34,7 +34,7 @@ describe('migrateSubblockIds', () => {
         }),
       }
 
-      const migrated = migrateSubblockIds(blocks)
+      const { blocks, migrated } = migrateSubblockIds(input)
 
       expect(migrated).toBe(true)
       expect(blocks['b1'].subBlocks['knowledgeBaseSelector']).toEqual({
@@ -47,7 +47,7 @@ describe('migrateSubblockIds', () => {
     })
 
     it('should prefer new key when both old and new exist', () => {
-      const blocks: Record<string, BlockState> = {
+      const input: Record<string, BlockState> = {
         b1: makeBlock({
           type: 'knowledge',
           subBlocks: {
@@ -65,15 +65,15 @@ describe('migrateSubblockIds', () => {
         }),
       }
 
-      const migrated = migrateSubblockIds(blocks)
+      const { blocks, migrated } = migrateSubblockIds(input)
 
       expect(migrated).toBe(true)
       expect(blocks['b1'].subBlocks['knowledgeBaseSelector'].value).toBe('fresh-kb')
       expect(blocks['b1'].subBlocks['knowledgeBaseId']).toBeUndefined()
     })
 
     it('should not touch blocks that already use the new key', () => {
-      const blocks: Record<string, BlockState> = {
+      const input: Record<string, BlockState> = {
         b1: makeBlock({
           type: 'knowledge',
           subBlocks: {
@@ -86,15 +86,36 @@ describe('migrateSubblockIds', () => {
         }),
       }
 
-      const migrated = migrateSubblockIds(blocks)
+      const { blocks, migrated } = migrateSubblockIds(input)
 
       expect(migrated).toBe(false)
       expect(blocks['b1'].subBlocks['knowledgeBaseSelector'].value).toBe('kb-uuid')
     })
   })
 
+  it('should not mutate the input blocks', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        subBlocks: {
+          knowledgeBaseId: {
+            id: 'knowledgeBaseId',
+            type: 'knowledge-base-selector',
+            value: 'kb-uuid',
+          },
+        },
+      }),
+    }
+
+    const { blocks } = migrateSubblockIds(input)
+
+    expect(input['b1'].subBlocks['knowledgeBaseId']).toBeDefined()
+    expect(blocks['b1'].subBlocks['knowledgeBaseSelector']).toBeDefined()
+    expect(blocks).not.toBe(input)
+  })
+
   it('should skip blocks with no registered migrations', () => {
-    const blocks: Record<string, BlockState> = {
+    const input: Record<string, BlockState> = {
       b1: makeBlock({
         type: 'function',
         subBlocks: {
@@ -103,14 +124,14 @@ describe('migrateSubblockIds', () => {
       }),
     }
 
-    const migrated = migrateSubblockIds(blocks)
+    const { blocks, migrated } = migrateSubblockIds(input)
 
     expect(migrated).toBe(false)
     expect(blocks['b1'].subBlocks['code'].value).toBe('console.log("hi")')
   })
 
   it('should migrate multiple blocks in one pass', () => {
-    const blocks: Record<string, BlockState> = {
+    const input: Record<string, BlockState> = {
       b1: makeBlock({
         id: 'b1',
         type: 'knowledge',
@@ -142,7 +163,7 @@ describe('migrateSubblockIds', () => {
       }),
     }
 
-    const migrated = migrateSubblockIds(blocks)
+    const { blocks, migrated } = migrateSubblockIds(input)
 
     expect(migrated).toBe(true)
     expect(blocks['b1'].subBlocks['knowledgeBaseSelector'].value).toBe('kb-1')
@@ -151,11 +172,11 @@ describe('migrateSubblockIds', () => {
   })
 
   it('should handle blocks with empty subBlocks', () => {
-    const blocks: Record<string, BlockState> = {
+    const input: Record<string, BlockState> = {
       b1: makeBlock({ type: 'knowledge', subBlocks: {} }),
     }
 
-    const migrated = migrateSubblockIds(blocks)
+    const { migrated } = migrateSubblockIds(input)
 
     expect(migrated).toBe(false)
   })
diff --git a/apps/sim/lib/workflows/migrations/subblock-migrations.ts b/apps/sim/lib/workflows/migrations/subblock-migrations.ts
@@ -20,58 +20,71 @@ export const SUBBLOCK_ID_MIGRATIONS: Record<string, Record<string, string>> = {
 
 /**
  * Migrates legacy subblock IDs inside a single block's subBlocks map.
- * If an old key is found and the new key does not already exist, the entry
- * is moved to the new key. When both exist the new key wins (user already
- * re-saved) and the old entry is removed to avoid orphans.
- *
- * Returns true if any migration was applied.
+ * Returns a new subBlocks record if anything changed, or the original if not.
  */
-function migrateBlockSubblockIds(block: BlockState, renames: Record<string, string>): boolean {
-  const subBlocks = block.subBlocks
-  if (!subBlocks) return false
-
+function migrateBlockSubblockIds(
+  subBlocks: Record<string, BlockState['subBlocks'][string]>,
+  renames: Record<string, string>
+): { subBlocks: Record<string, BlockState['subBlocks'][string]>; migrated: boolean } {
   let migrated = false
 
+  for (const oldId of Object.keys(renames)) {
+    if (oldId in subBlocks) {
+      migrated = true
+      break
+    }
+  }
+
+  if (!migrated) return { subBlocks, migrated: false }
+
+  const result = { ...subBlocks }
+
   for (const [oldId, newId] of Object.entries(renames)) {
-    if (!(oldId in subBlocks)) continue
+    if (!(oldId in result)) continue
 
-    if (newId in subBlocks) {
-      delete subBlocks[oldId]
-      migrated = true
+    if (newId in result) {
+      delete result[oldId]
       continue
     }
 
-    const oldEntry = subBlocks[oldId]
-    subBlocks[newId] = { ...oldEntry, id: newId }
-    delete subBlocks[oldId]
-    migrated = true
+    const oldEntry = result[oldId]
+    result[newId] = { ...oldEntry, id: newId }
+    delete result[oldId]
   }
 
-  return migrated
+  return { subBlocks: result, migrated: true }
 }
 
 /**
  * Applies subblock-ID migrations to every block in a workflow.
- * Safe to call on any state – blocks whose type has no registered
- * migrations are left untouched.
- *
- * Mutates `blocks` in place and returns whether anything changed.
+ * Returns a new blocks record with migrated subBlocks where needed.
  */
-export function migrateSubblockIds(blocks: Record<string, BlockState>): boolean {
+export function migrateSubblockIds(blocks: Record<string, BlockState>): {
+  blocks: Record<string, BlockState>
+  migrated: boolean
+} {
   let anyMigrated = false
+  const result: Record<string, BlockState> = {}
 
-  for (const block of Object.values(blocks)) {
+  for (const [blockId, block] of Object.entries(blocks)) {
     const renames = SUBBLOCK_ID_MIGRATIONS[block.type]
-    if (!renames) continue
+    if (!renames || !block.subBlocks) {
+      result[blockId] = block
+      continue
+    }
 
-    if (migrateBlockSubblockIds(block, renames)) {
+    const { subBlocks, migrated } = migrateBlockSubblockIds(block.subBlocks, renames)
+    if (migrated) {
       logger.info('Migrated legacy subblock IDs', {
         blockId: block.id,
         blockType: block.type,
       })
       anyMigrated = true
+      result[blockId] = { ...block, subBlocks }
+    } else {
+      result[blockId] = block
     }
   }
 
-  return anyMigrated
+  return { blocks: result, migrated: anyMigrated }
 }
diff --git a/apps/sim/lib/workflows/persistence/utils.ts b/apps/sim/lib/workflows/persistence/utils.ts
@@ -114,15 +114,13 @@ export async function loadDeployedWorkflowState(
       resolvedWorkspaceId = wfRow?.workspaceId ?? undefined
     }
 
-    const resolvedBlocks = state.blocks || {}
-    const { blocks: credMigratedBlocks } = resolvedWorkspaceId
-      ? await migrateCredentialIds(resolvedBlocks, resolvedWorkspaceId)
-      : { blocks: resolvedBlocks }
-
-    migrateSubblockIds(credMigratedBlocks)
+    const { blocks: migratedBlocks } = await applyBlockMigrations(
+      state.blocks || {},
+      resolvedWorkspaceId
+    )
 
     return {
-      blocks: credMigratedBlocks,
+      blocks: migratedBlocks,
       edges: state.edges || [],
       loops: state.loops || {},
       parallels: state.parallels || {},
@@ -136,6 +134,50 @@ export async function loadDeployedWorkflowState(
   }
 }
 
+interface MigrationContext {
+  blocks: Record<string, BlockState>
+  workspaceId?: string
+  migrated: boolean
+}
+
+type BlockMigration = (ctx: MigrationContext) => MigrationContext | Promise<MigrationContext>
+
+function createMigrationPipeline(migrations: BlockMigration[]) {
+  return async (
+    blocks: Record<string, BlockState>,
+    workspaceId?: string
+  ): Promise<{ blocks: Record<string, BlockState>; migrated: boolean }> => {
+    let ctx: MigrationContext = { blocks, workspaceId, migrated: false }
+    for (const migration of migrations) {
+      ctx = await migration(ctx)
+    }
+    return { blocks: ctx.blocks, migrated: ctx.migrated }
+  }
+}
+
+const applyBlockMigrations = createMigrationPipeline([
+  (ctx) => {
+    const { blocks } = sanitizeAgentToolsInBlocks(ctx.blocks)
+    return { ...ctx, blocks }
+  },
+
+  (ctx) => ({
+    ...ctx,
+    blocks: migrateAgentBlocksToMessagesFormat(ctx.blocks),
+  }),
+
+  async (ctx) => {
+    if (!ctx.workspaceId) return ctx
+    const { blocks, migrated } = await migrateCredentialIds(ctx.blocks, ctx.workspaceId)
+    return { ...ctx, blocks, migrated: ctx.migrated || migrated }
+  },
+
+  (ctx) => {
+    const { blocks, migrated } = migrateSubblockIds(ctx.blocks)
+    return { ...ctx, blocks, migrated: ctx.migrated || migrated }
+  },
+])
+
 /**
  * Migrates agent blocks from old format (systemPrompt/userPrompt) to new format (messages array)
  * This ensures backward compatibility for workflows created before the messages-input refactor.
@@ -359,22 +401,16 @@ export async function loadWorkflowFromNormalizedTables(
       blocksMap[block.id] = assembled
     })
 
-    // Sanitize any invalid custom tools in agent blocks to prevent client crashes
-    const { blocks: sanitizedBlocks } = sanitizeAgentToolsInBlocks(blocksMap)
-
-    // Migrate old agent block format (systemPrompt/userPrompt) to new messages array format
-    const migratedBlocks = migrateAgentBlocksToMessagesFormat(sanitizedBlocks)
-
-    // Migrate legacy account.id → credential.id in OAuth subblocks
-    const { blocks: credMigratedBlocks, migrated: credentialsMigrated } = workflowRow?.workspaceId
-      ? await migrateCredentialIds(migratedBlocks, workflowRow.workspaceId)
-      : { blocks: migratedBlocks, migrated: false }
+    const { blocks: finalBlocks, migrated } = await applyBlockMigrations(
+      blocksMap,
+      workflowRow?.workspaceId ?? undefined
+    )
 
-    if (credentialsMigrated) {
+    if (migrated) {
       Promise.resolve().then(async () => {
         try {
-          for (const [blockId, block] of Object.entries(credMigratedBlocks)) {
-            if (block.subBlocks !== migratedBlocks[blockId]?.subBlocks) {
+          for (const [blockId, block] of Object.entries(finalBlocks)) {
+            if (block.subBlocks !== blocksMap[blockId]?.subBlocks) {
               await db
                 .update(workflowBlocks)
                 .set({ subBlocks: block.subBlocks, updatedAt: new Date() })
@@ -384,23 +420,7 @@ export async function loadWorkflowFromNormalizedTables(
             }
           }
         } catch (err) {
-          logger.warn('Failed to persist credential ID migration', { workflowId, error: err })
-        }
-      })
-    }
-
-    const subblockMigrated = migrateSubblockIds(credMigratedBlocks)
-    if (subblockMigrated) {
-      Promise.resolve().then(async () => {
-        try {
-          for (const [blockId, block] of Object.entries(credMigratedBlocks)) {
-            await db
-              .update(workflowBlocks)
-              .set({ subBlocks: block.subBlocks, updatedAt: new Date() })
-              .where(and(eq(workflowBlocks.id, blockId), eq(workflowBlocks.workflowId, workflowId)))
-          }
-        } catch (err) {
-          logger.warn('Failed to persist subblock ID migration', { workflowId, error: err })
+          logger.warn('Failed to persist block migrations', { workflowId, error: err })
         }
       })
     }
@@ -441,13 +461,13 @@ export async function loadWorkflowFromNormalizedTables(
           forEachItems: (config as Loop).forEachItems ?? '',
           whileCondition: (config as Loop).whileCondition ?? '',
           doWhileCondition: (config as Loop).doWhileCondition ?? '',
-          enabled: credMigratedBlocks[subflow.id]?.enabled ?? true,
+          enabled: finalBlocks[subflow.id]?.enabled ?? true,
         }
         loops[subflow.id] = loop
 
-        if (credMigratedBlocks[subflow.id]) {
-          const block = credMigratedBlocks[subflow.id]
-          credMigratedBlocks[subflow.id] = {
+        if (finalBlocks[subflow.id]) {
+          const block = finalBlocks[subflow.id]
+          finalBlocks[subflow.id] = {
             ...block,
             data: {
               ...block.data,
@@ -468,7 +488,7 @@ export async function loadWorkflowFromNormalizedTables(
             (config as Parallel).parallelType === 'collection'
               ? (config as Parallel).parallelType
               : 'count',
-          enabled: credMigratedBlocks[subflow.id]?.enabled ?? true,
+          enabled: finalBlocks[subflow.id]?.enabled ?? true,
         }
         parallels[subflow.id] = parallel
       } else {
@@ -477,7 +497,7 @@ export async function loadWorkflowFromNormalizedTables(
     })
 
     return {
-      blocks: credMigratedBlocks,
+      blocks: finalBlocks,
       edges: edgesArray,
       loops,
       parallels,
