improvement(canonical): backfill for canonical modes on config changes (#3447)

* improvement(canonical): backfill for canonical modes on config changes

* persist data changes to db

apps/sim/lib/workflows/migrations/subblock-migrations.test.ts

@@ -1,9 +1,12 @@
 /**
  * @vitest-environment node
  */
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 import type { BlockState } from '@/stores/workflows/workflow/types'
-import { migrateSubblockIds } from './subblock-migrations'
+
+vi.unmock('@/blocks/registry')
+
+import { backfillCanonicalModes, migrateSubblockIds } from './subblock-migrations'
 
 function makeBlock(overrides: Partial<BlockState> & { type: string }): BlockState {
   return {
@@ -181,3 +184,185 @@ describe('migrateSubblockIds', () => {
     expect(migrated).toBe(false)
   })
 })
+
+describe('backfillCanonicalModes', () => {
+  it('should add missing canonicalModes entry for knowledge block with basic value', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: {},
+        subBlocks: {
+          operation: { id: 'operation', type: 'dropdown', value: 'search' },
+          knowledgeBaseSelector: {
+            id: 'knowledgeBaseSelector',
+            type: 'knowledge-base-selector',
+            value: 'kb-uuid',
+          },
+        },
+      }),
+    }
+
+    const { blocks, migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(true)
+    const modes = blocks.b1.data?.canonicalModes as Record<string, string>
+    expect(modes.knowledgeBaseId).toBe('basic')
+  })
+
+  it('should resolve to advanced when only the advanced value is set', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: {},
+        subBlocks: {
+          operation: { id: 'operation', type: 'dropdown', value: 'search' },
+          manualKnowledgeBaseId: {
+            id: 'manualKnowledgeBaseId',
+            type: 'short-input',
+            value: 'kb-uuid-manual',
+          },
+        },
+      }),
+    }
+
+    const { blocks, migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(true)
+    const modes = blocks.b1.data?.canonicalModes as Record<string, string>
+    expect(modes.knowledgeBaseId).toBe('advanced')
+  })
+
+  it('should not overwrite existing canonicalModes entries', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: { canonicalModes: { knowledgeBaseId: 'advanced' } },
+        subBlocks: {
+          knowledgeBaseSelector: {
+            id: 'knowledgeBaseSelector',
+            type: 'knowledge-base-selector',
+            value: 'kb-uuid',
+          },
+        },
+      }),
+    }
+
+    const { blocks, migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(false)
+    const modes = blocks.b1.data?.canonicalModes as Record<string, string>
+    expect(modes.knowledgeBaseId).toBe('advanced')
+  })
+
+  it('should skip blocks with no canonical pairs in their config', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'function',
+        data: {},
+        subBlocks: {
+          code: { id: 'code', type: 'code', value: '' },
+        },
+      }),
+    }
+
+    const { migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(false)
+  })
+
+  it('should not mutate the input blocks', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: {},
+        subBlocks: {
+          knowledgeBaseSelector: {
+            id: 'knowledgeBaseSelector',
+            type: 'knowledge-base-selector',
+            value: 'kb-uuid',
+          },
+        },
+      }),
+    }
+
+    const { blocks } = backfillCanonicalModes(input)
+
+    expect(input.b1.data?.canonicalModes).toBeUndefined()
+    expect((blocks.b1.data?.canonicalModes as Record<string, string>).knowledgeBaseId).toBe('basic')
+    expect(blocks).not.toBe(input)
+  })
+
+  it('should resolve correctly when existing field became the basic variant', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: {},
+        subBlocks: {
+          operation: { id: 'operation', type: 'dropdown', value: 'search' },
+          knowledgeBaseSelector: {
+            id: 'knowledgeBaseSelector',
+            type: 'knowledge-base-selector',
+            value: 'kb-uuid',
+          },
+          manualKnowledgeBaseId: {
+            id: 'manualKnowledgeBaseId',
+            type: 'short-input',
+            value: '',
+          },
+        },
+      }),
+    }
+
+    const { blocks, migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(true)
+    const modes = blocks.b1.data?.canonicalModes as Record<string, string>
+    expect(modes.knowledgeBaseId).toBe('basic')
+  })
+
+  it('should resolve correctly when existing field became the advanced variant', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: {},
+        subBlocks: {
+          operation: { id: 'operation', type: 'dropdown', value: 'search' },
+          knowledgeBaseSelector: {
+            id: 'knowledgeBaseSelector',
+            type: 'knowledge-base-selector',
+            value: '',
+          },
+          manualKnowledgeBaseId: {
+            id: 'manualKnowledgeBaseId',
+            type: 'short-input',
+            value: 'manually-entered-kb-id',
+          },
+        },
+      }),
+    }
+
+    const { blocks, migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(true)
+    const modes = blocks.b1.data?.canonicalModes as Record<string, string>
+    expect(modes.knowledgeBaseId).toBe('advanced')
+  })
+
+  it('should default to basic when neither value is set', () => {
+    const input: Record<string, BlockState> = {
+      b1: makeBlock({
+        type: 'knowledge',
+        data: {},
+        subBlocks: {
+          operation: { id: 'operation', type: 'dropdown', value: 'search' },
+        },
+      }),
+    }
+
+    const { blocks, migrated } = backfillCanonicalModes(input)
+
+    expect(migrated).toBe(true)
+    const modes = blocks.b1.data?.canonicalModes as Record<string, string>
+    expect(modes.knowledgeBaseId).toBe('basic')
+  })
+})

apps/sim/lib/workflows/migrations/subblock-migrations.ts

@@ -1,4 +1,11 @@
 import { createLogger } from '@sim/logger'
+import {
+  buildCanonicalIndex,
+  buildSubBlockValues,
+  isCanonicalPair,
+  resolveCanonicalMode,
+} from '@/lib/workflows/subblocks/visibility'
+import { getBlock } from '@/blocks'
 import type { BlockState } from '@/stores/workflows/workflow/types'
 
 const logger = createLogger('SubblockMigrations')
@@ -88,3 +95,59 @@ export function migrateSubblockIds(blocks: Record<string, BlockState>): {
 
   return { blocks: result, migrated: anyMigrated }
 }
+
+/**
+ * Backfills missing `canonicalModes` entries in block data.
+ *
+ * When a canonical pair is added to a block definition, existing blocks
+ * won't have the entry in `data.canonicalModes`. Without it the editor
+ * toggle may not render correctly. This resolves the correct mode based
+ * on which subblock value is populated and adds the missing entry.
+ */
+export function backfillCanonicalModes(blocks: Record<string, BlockState>): {
+  blocks: Record<string, BlockState>
+  migrated: boolean
+} {
+  let anyMigrated = false
+  const result: Record<string, BlockState> = {}
+
+  for (const [blockId, block] of Object.entries(blocks)) {
+    const blockConfig = getBlock(block.type)
+    if (!blockConfig?.subBlocks || !block.subBlocks) {
+      result[blockId] = block
+      continue
+    }
+
+    const canonicalIndex = buildCanonicalIndex(blockConfig.subBlocks)
+    const pairs = Object.values(canonicalIndex.groupsById).filter(isCanonicalPair)
+    if (pairs.length === 0) {
+      result[blockId] = block
+      continue
+    }
+
+    const existing = (block.data?.canonicalModes ?? {}) as Record<string, 'basic' | 'advanced'>
+    let patched: Record<string, 'basic' | 'advanced'> | null = null
+
+    const values = buildSubBlockValues(block.subBlocks)
+
+    for (const group of pairs) {
+      if (existing[group.canonicalId] != null) continue
+
+      const resolved = resolveCanonicalMode(group, values)
+      if (!patched) patched = { ...existing }
+      patched[group.canonicalId] = resolved
+    }
+
+    if (patched) {
+      anyMigrated = true
+      result[blockId] = {
+        ...block,
+        data: { ...(block.data ?? {}), canonicalModes: patched },
+      }
+    } else {
+      result[blockId] = block
+    }
+  }
+
+  return { blocks: result, migrated: anyMigrated }
+}

apps/sim/lib/workflows/persistence/utils.ts

@@ -14,7 +14,10 @@ import { and, desc, eq, inArray, sql } from 'drizzle-orm'
 import type { Edge } from 'reactflow'
 import { v4 as uuidv4 } from 'uuid'
 import type { DbOrTx } from '@/lib/db/types'
-import { migrateSubblockIds } from '@/lib/workflows/migrations/subblock-migrations'
+import {
+  backfillCanonicalModes,
+  migrateSubblockIds,
+} from '@/lib/workflows/migrations/subblock-migrations'
 import { sanitizeAgentToolsInBlocks } from '@/lib/workflows/sanitization/validation'
 import type { BlockState, Loop, Parallel, WorkflowState } from '@/stores/workflows/workflow/types'
 import { SUBFLOW_TYPES } from '@/stores/workflows/workflow/types'
@@ -176,6 +179,11 @@ const applyBlockMigrations = createMigrationPipeline([
     const { blocks, migrated } = migrateSubblockIds(ctx.blocks)
     return { ...ctx, blocks, migrated: ctx.migrated || migrated }
   },
+
+  (ctx) => {
+    const { blocks, migrated } = backfillCanonicalModes(ctx.blocks)
+    return { ...ctx, blocks, migrated: ctx.migrated || migrated }
+  },
 ])
 
 /**
@@ -410,10 +418,14 @@ export async function loadWorkflowFromNormalizedTables(
       Promise.resolve().then(async () => {
         try {
           for (const [blockId, block] of Object.entries(finalBlocks)) {
-            if (block.subBlocks !== blocksMap[blockId]?.subBlocks) {
+            if (block !== blocksMap[blockId]) {
               await db
                 .update(workflowBlocks)
-                .set({ subBlocks: block.subBlocks, updatedAt: new Date() })
+                .set({
+                  subBlocks: block.subBlocks,
+                  data: block.data,
+                  updatedAt: new Date(),
+                })
                 .where(
                   and(eq(workflowBlocks.id, blockId), eq(workflowBlocks.workflowId, workflowId))
                 )