fix(agent): lock external saves and distinguish workflow targets

Author: j15z

apps/realtime/src/database/agent-tool-permissions.test.ts

@@ -8,7 +8,7 @@ const tool = {
   usageControlExpression: 'force',
 }
 function block(
-  tools = [tool],
+  tools: Record<string, unknown>[] = [tool],
   modes: Record<string, string> = { '0:agentToolUsageControl': 'advanced' }
 ) {
   return {
@@ -89,4 +89,25 @@ describe('variable permission changes', () => {
       )
     ).toBe(true)
   })
+
+  it('rejects transferring a variable permission to a different workflow target', () => {
+    const workflowTool = {
+      ...tool,
+      type: 'workflow_input',
+      toolId: 'workflow_executor',
+      params: { ...tool.params, workflowId: 'workflow-a' },
+    }
+    const before = block([workflowTool])
+    expect(hasAgentToolPermissionChanges([structuredClone(before)], [before])).toBe(false)
+    expect(
+      hasAgentToolPermissionChanges(
+        [
+          block([
+            { ...workflowTool, params: { ...workflowTool.params, workflowId: 'workflow-b' } },
+          ]),
+        ],
+        [before]
+      )
+    ).toBe(true)
+  })
 })

apps/sim/lib/workflows/persistence/persist-agent-tool-permission-mode.test.ts

@@ -3,8 +3,8 @@
  */
 import {
   createAgentBlock,
-  databaseMock,
   dbChainMock,
+  dbChainMockFns,
   queueTableRows,
   resetDbChainMock,
   schemaMock,
@@ -17,8 +17,6 @@ const mocks = vi.hoisted(() => ({
   saveRaw: vi.fn(),
 }))
 
-vi.mock('@sim/db', () => ({ ...databaseMock, workflow: schemaMock.workflow }))
-
 vi.mock('@/lib/core/config/feature-flags', () => ({
   isFeatureEnabled: mocks.isFeatureEnabled,
 }))
@@ -181,4 +179,28 @@ describe('agent tool Permission Mode write gate', () => {
     ).rejects.toMatchObject({ code: 'validation' })
     expect(mocks.saveRaw).not.toHaveBeenCalled()
   })
+
+  it('waits for the caller transaction lock before reading retained permissions', async () => {
+    const state = stateWithTool(
+      { type: 'function', usageControlExpression: 'force' },
+      { '0:agentToolUsageControl': 'advanced' }
+    )
+    let releaseLock!: (rows: { id: string }[]) => void
+    const lock = new Promise<{ id: string }[]>((resolve) => {
+      releaseLock = resolve
+    })
+    dbChainMockFns.for.mockReturnValueOnce(lock)
+
+    const save = saveWorkflowToNormalizedTables('workflow-1', state, GOVERNANCE, dbChainMock.db)
+    await vi.waitFor(() => expect(dbChainMockFns.for).toHaveBeenCalledWith('update'), {
+      interval: 1,
+    })
+    expect(mocks.isFeatureEnabled).not.toHaveBeenCalled()
+    expect(mocks.saveRaw).not.toHaveBeenCalled()
+
+    queueTableRows(schemaMock.workflowBlocks, [])
+    releaseLock([{ id: 'workflow-1' }])
+    await expect(save).rejects.toMatchObject({ code: 'validation' })
+    expect(mocks.saveRaw).not.toHaveBeenCalled()
+  })
 })

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

@@ -1,11 +1,10 @@
 import {
   db,
   runOutsideTransactionContext,
-  workflow,
   workflowDeploymentOperation,
   workflowDeploymentVersion,
 } from '@sim/db'
-import { credential } from '@sim/db/schema'
+import { credential, workflow } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getActiveWorkflowContext } from '@sim/platform-authz/workflow'
 import { getErrorMessage } from '@sim/utils/errors'
@@ -673,6 +672,7 @@ export async function saveWorkflowToNormalizedTables(
   externalTx?: DbOrTx
 ): Promise<{ success: boolean; error?: string }> {
   await assertNoWithheldBlockType(governance, Object.values(state.blocks))
+  if (externalTx) await lockWorkflowForUpdate(externalTx, workflowId)
   await assertAgentToolPermissionModeEnabled(Object.values(state.blocks), {
     workflowId,
     tx: externalTx,

packages/workflow-types/src/agent-tool-permissions.ts

@@ -62,6 +62,7 @@ function toolPermissions(block: AgentToolPermissionBlock): ToolPermission[] {
           fn.name,
           params.serverId,
           params.toolName,
+          params.workflowId,
         ]),
         expression: tool.usageControlExpression,
         advanced,