fix(agent): keep nested tool basic/advanced modes attached to their tools

apps/realtime/src/database/operations.test.ts

@@ -121,3 +121,51 @@ describe('search replacement persistence', () => {
     expect(mockSet).toHaveBeenCalledTimes(1)
   })
 })
+
+describe('subblock update with canonical modes persistence', () => {
+  const tools = [{ type: 'jira', params: { manualProjectId: '{{PROJECT}}' } }]
+  const canonicalModes = { '0:projectId': 'advanced' as const, model: 'basic' as const }
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockTransaction.mockImplementation(
+      async (callback: (tx: typeof transaction) => Promise<void>) => callback(transaction)
+    )
+    mockSet.mockReturnValue({ where: vi.fn().mockResolvedValue(undefined) })
+  })
+
+  function updateTools(block: Record<string, unknown>) {
+    mockSelectWhere.mockResolvedValue([
+      {
+        id: 'agent-1',
+        locked: false,
+        data: { width: 350, canonicalModes: { '1:projectId': 'advanced' } },
+        subBlocks: { tools: { id: 'tools', type: 'tool-input', value: [] } },
+        ...block,
+      },
+    ])
+    return persistWorkflowOperation('workflow-1', {
+      operation: SUBBLOCK_OPERATIONS.UPDATE_WITH_CANONICAL_MODES,
+      target: OPERATION_TARGETS.SUBBLOCK,
+      timestamp: Date.now(),
+      payload: { blockId: 'agent-1', subblockId: 'tools', value: tools, canonicalModes },
+    })
+  }
+
+  it('writes the subblock value and replaces canonical modes in one block update', async () => {
+    await expect(updateTools({})).resolves.toBeUndefined()
+
+    expect(mockSet).toHaveBeenCalledTimes(2)
+    expect(mockSet).toHaveBeenLastCalledWith(
+      expect.objectContaining({
+        subBlocks: { tools: { id: 'tools', type: 'tool-input', value: tools } },
+        data: { width: 350, canonicalModes },
+      })
+    )
+  })
+
+  it('rejects a locked block without writing either field', async () => {
+    await expect(updateTools({ locked: true })).rejects.toThrow('is locked')
+    expect(mockSet).toHaveBeenCalledTimes(1)
+  })
+})

apps/realtime/src/database/operations.ts

@@ -26,6 +26,7 @@ import {
 import { randomFloat } from '@sim/utils/random'
 import { loadWorkflowFromNormalizedTablesRaw } from '@sim/workflow-persistence/load'
 import { mergeSubBlockValues } from '@sim/workflow-persistence/subblocks'
+import type { DbOrTx } from '@sim/workflow-persistence/types'
 import {
   filterAcyclicEdges,
   filterUniqueWorkflowEdges,
@@ -1989,6 +1990,38 @@ async function handleSubflowOperationTx(
   }
 }
 
+/** Every block in the workflow by id, for the locked-container check subblock writes need. */
+async function loadSubblockUpdateBlocks(tx: DbOrTx, workflowId: string) {
+  const allBlocks = await tx
+    .select({
+      id: workflowBlocks.id,
+      subBlocks: workflowBlocks.subBlocks,
+      locked: workflowBlocks.locked,
+      data: workflowBlocks.data,
+    })
+    .from(workflowBlocks)
+    .where(eq(workflowBlocks.workflowId, workflowId))
+  return Object.fromEntries(allBlocks.map((block) => [block.id, block]))
+}
+
+/**
+ * The block a subblock write targets, rejecting one that is missing, locked, or in a locked
+ * container.
+ */
+function getWritableSubblockUpdateBlock(
+  blocksById: Awaited<ReturnType<typeof loadSubblockUpdateBlocks>>,
+  blockId: string
+) {
+  const block = blocksById[blockId]
+  if (!block) {
+    throw new Error(`Block ${blockId} not found`)
+  }
+  if (isWorkflowBlockProtected(blockId, blocksById)) {
+    throw new Error(`Block ${blockId} is locked or inside a locked container`)
+  }
+  return block
+}
+
 // Subblock operations - targeted value updates without replacing workflow state
 async function handleSubblockOperationTx(
   tx: any,
@@ -2003,35 +2036,15 @@ async function handleSubblockOperationTx(
         return
       }
 
-      const allBlocks = await tx
-        .select({
-          id: workflowBlocks.id,
-          subBlocks: workflowBlocks.subBlocks,
-          locked: workflowBlocks.locked,
-          data: workflowBlocks.data,
-        })
-        .from(workflowBlocks)
-        .where(eq(workflowBlocks.workflowId, workflowId))
-
-      type SubblockUpdateBlockRecord = (typeof allBlocks)[number]
-      const blocksById: Record<string, SubblockUpdateBlockRecord> = Object.fromEntries(
-        allBlocks.map((block: SubblockUpdateBlockRecord) => [block.id, block])
-      )
+      const blocksById = await loadSubblockUpdateBlocks(tx, workflowId)
 
       for (const update of updates) {
         const { blockId, subblockId, value, expectedValue } = update
         if (!blockId || !subblockId) {
           throw new Error('Missing required fields for subblock batch update')
         }
 
-        const block = blocksById[blockId]
-        if (!block) {
-          throw new Error(`Block ${blockId} not found`)
-        }
-
-        if (isWorkflowBlockProtected(blockId, blocksById)) {
-          throw new Error(`Block ${blockId} is locked or inside a locked container`)
-        }
+        const block = getWritableSubblockUpdateBlock(blocksById, blockId)
 
         const subBlocks = { ...((block.subBlocks as Record<string, any>) || {}) }
         const currentSubBlock = subBlocks[subblockId]
@@ -2060,6 +2073,36 @@ async function handleSubblockOperationTx(
       break
     }
 
+    case SUBBLOCK_OPERATIONS.UPDATE_WITH_CANONICAL_MODES: {
+      const { blockId, subblockId, value, canonicalModes } = payload
+      if (!blockId || !subblockId || !canonicalModes) {
+        throw new Error('Missing required fields for subblock update with canonical modes')
+      }
+
+      const blocksById = await loadSubblockUpdateBlocks(tx, workflowId)
+      const block = getWritableSubblockUpdateBlock(blocksById, blockId)
+
+      const subBlocks = {
+        ...((block.subBlocks as Record<string, Record<string, unknown>> | null) || {}),
+      }
+      const currentSubBlock = subBlocks[subblockId]
+      subBlocks[subblockId] = currentSubBlock
+        ? { ...currentSubBlock, value }
+        : { id: subblockId, type: 'unknown', value }
+
+      await tx
+        .update(workflowBlocks)
+        .set({
+          subBlocks,
+          data: { ...((block.data as Record<string, unknown>) || {}), canonicalModes },
+          updatedAt: new Date(),
+        })
+        .where(and(eq(workflowBlocks.id, blockId), eq(workflowBlocks.workflowId, workflowId)))
+
+      logger.debug(`Updated subblock ${blockId}.${subblockId} with canonical modes`)
+      break
+    }
+
     default:
       logger.warn(`Unknown subblock operation: ${operation}`)
       throw new Error(`Unsupported subblock operation: ${operation}`)

apps/realtime/src/middleware/permissions.test.ts

@@ -114,6 +114,11 @@ describe('checkRolePermission', () => {
       const result = checkRolePermission('write', 'subblock-batch-update')
       expectPermissionAllowed(result)
     })
+
+    it('should allow subblock-update-with-canonical-modes operation', () => {
+      const result = checkRolePermission('write', 'subblock-update-with-canonical-modes')
+      expectPermissionAllowed(result)
+    })
   })
 
   describe('read role', () => {
@@ -155,6 +160,11 @@ describe('checkRolePermission', () => {
       expectPermissionDenied(result, 'read')
     })
 
+    it('should deny subblock-update-with-canonical-modes operation for read role', () => {
+      const result = checkRolePermission('read', 'subblock-update-with-canonical-modes')
+      expectPermissionDenied(result, 'read')
+    })
+
     it('should deny toggle-enabled operation for read role', () => {
       const result = checkRolePermission('read', 'toggle-enabled')
       expectPermissionDenied(result, 'read')

apps/realtime/src/middleware/permissions.ts

@@ -52,6 +52,7 @@ const WRITE_OPERATIONS: string[] = [
   // Subblock operations
   SUBBLOCK_OPERATIONS.UPDATE,
   SUBBLOCK_OPERATIONS.BATCH_UPDATE,
+  SUBBLOCK_OPERATIONS.UPDATE_WITH_CANONICAL_MODES,
   // Variable operations
   VARIABLE_OPERATIONS.UPDATE,
   // Workflow operations

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx

@@ -392,15 +392,8 @@ export const ToolInput = memo(function ToolInput({
       [blockId]
     )
   )
-  const { collaborativeSetBlockCanonicalMode, collaborativeSetBlockCanonicalModes } =
+  const { collaborativeSetBlockCanonicalMode, collaborativeSetSubblockValueWithCanonicalModes } =
     useCollaborativeWorkflow()
-  const reindexCanonicalModesOnMutate = useCallback(
-    (oldTools: StoredTool[], newTools: StoredTool[]) => {
-      const next = reindexToolCanonicalModes(oldTools, newTools, canonicalModeOverrides)
-      if (next) collaborativeSetBlockCanonicalModes(blockId, next)
-    },
-    [canonicalModeOverrides, collaborativeSetBlockCanonicalModes, blockId]
-  )
 
   const value = isPreview ? previewValue : storeValue
 
@@ -412,6 +405,39 @@ export const ToolInput = memo(function ToolInput({
       ? (value as StoredTool[])
       : []
 
+  /**
+   * Commits a tool list that moves or drops selected tools. Their canonical-mode overrides are
+   * keyed by position, so when any must move they persist in the same operation as the list.
+   * `positionedTools` is the list holding the kept tool references when `nextTools` clones them.
+   */
+  const setToolsWithReindexedModes = useCallback(
+    (nextTools: StoredTool[], positionedTools: StoredTool[] = nextTools) => {
+      const canonicalModes = reindexToolCanonicalModes(
+        selectedTools,
+        positionedTools,
+        canonicalModeOverrides
+      )
+      if (!canonicalModes) {
+        setStoreValue(nextTools)
+        return
+      }
+      collaborativeSetSubblockValueWithCanonicalModes(
+        blockId,
+        subBlockId,
+        structuredClone(nextTools),
+        canonicalModes
+      )
+    },
+    [
+      selectedTools,
+      canonicalModeOverrides,
+      setStoreValue,
+      collaborativeSetSubblockValueWithCanonicalModes,
+      blockId,
+      subBlockId,
+    ]
+  )
+
   // Tool categories the consuming block can't run (declared on its tool-input
   // subBlock): shown in the picker but greyed out with a tooltip instead of added.
   const blockType = useWorkflowStore(useCallback((state) => state.blocks[blockId]?.type, [blockId]))
@@ -857,10 +883,9 @@ export const ToolInput = memo(function ToolInput({
     (toolIndex: number) => {
       if (isPreview || disabled) return
       const updatedTools = selectedTools.filter((_, index) => index !== toolIndex)
-      reindexCanonicalModesOnMutate(selectedTools, updatedTools)
-      setStoreValue(updatedTools)
+      setToolsWithReindexedModes(updatedTools)
     },
-    [isPreview, disabled, selectedTools, reindexCanonicalModesOnMutate, setStoreValue]
+    [isPreview, disabled, selectedTools, setToolsWithReindexedModes]
   )
 
   const handleRemoveAllFromServer = useCallback(
@@ -869,10 +894,9 @@ export const ToolInput = memo(function ToolInput({
       const updatedTools = selectedTools.filter(
         (t) => !(t.type === 'mcp' && t.params?.serverId === serverId)
       )
-      reindexCanonicalModesOnMutate(selectedTools, updatedTools)
-      setStoreValue(updatedTools)
+      setToolsWithReindexedModes(updatedTools)
     },
-    [isPreview, disabled, selectedTools, reindexCanonicalModesOnMutate, setStoreValue]
+    [isPreview, disabled, selectedTools, setToolsWithReindexedModes]
   )
 
   const handleDeleteTool = useCallback(
@@ -900,11 +924,10 @@ export const ToolInput = memo(function ToolInput({
       })
 
       if (updatedTools.length !== selectedTools.length) {
-        reindexCanonicalModesOnMutate(selectedTools, updatedTools)
-        setStoreValue(updatedTools)
+        setToolsWithReindexedModes(updatedTools)
       }
     },
-    [selectedTools, customTools, reindexCanonicalModesOnMutate, setStoreValue]
+    [selectedTools, customTools, setToolsWithReindexedModes]
   )
 
   const handleParamChange = useCallback(
@@ -1077,8 +1100,7 @@ export const ToolInput = memo(function ToolInput({
       newTools.splice(adjustedDropIndex, 0, draggedTool)
     }
 
-    reindexCanonicalModesOnMutate(selectedTools, newTools)
-    setStoreValue(newTools)
+    setToolsWithReindexedModes(newTools)
     setDraggedIndex(null)
     setDragOverIndex(null)
   }
@@ -1177,8 +1199,7 @@ export const ToolInput = memo(function ToolInput({
               ...filteredTools.map((tool) => ({ ...tool, isExpanded: false })),
               serverBinding,
             ]
-            reindexCanonicalModesOnMutate(selectedTools, filteredTools)
-            setStoreValue(nextTools)
+            setToolsWithReindexedModes(nextTools, filteredTools)
             setMcpServerDrilldown(null)
             setOpen(false)
           },
@@ -1445,7 +1466,7 @@ export const ToolInput = memo(function ToolInput({
     supportsAdvancedMcpServer,
     availableWorkflows,
     isToolAlreadySelected,
-    reindexCanonicalModesOnMutate,
+    setToolsWithReindexedModes,
   ])
 
   return (