fix(agent): gate realtime tool permission writes

Author: j15z

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

@@ -1,5 +1,11 @@
 /** @vitest-environment node */
-import { OPERATION_TARGETS, SUBBLOCK_OPERATIONS } from '@sim/realtime-protocol/constants'
+import {
+  BLOCK_OPERATIONS,
+  BLOCKS_OPERATIONS,
+  OPERATION_TARGETS,
+  SUBBLOCK_OPERATIONS,
+  WORKFLOW_OPERATIONS,
+} from '@sim/realtime-protocol/constants'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const { mockTransaction, mockSelectWhere, mockSet } = vi.hoisted(() => ({
@@ -38,13 +44,28 @@ vi.mock('drizzle-orm', () => ({
 }))
 vi.mock('drizzle-orm/postgres-js', () => ({ drizzle: () => ({ transaction: mockTransaction }) }))
 vi.mock('postgres', () => ({ default: vi.fn() }))
-vi.mock('@/env', () => ({ env: { DATABASE_URL: 'postgres://localhost/test' } }))
+vi.mock('@/env', () => ({
+  env: { DATABASE_URL: 'postgres://localhost/test' },
+  getBaseUrl: () => 'http://localhost:3000',
+}))
 
+import { mergeSubBlockValues } from '@sim/workflow-persistence/subblocks'
+import { afterEach } from 'vitest'
 import { persistWorkflowOperation } from '@/database/operations'
 
+const mockFetch = vi.fn()
+beforeEach(() => {
+  vi.stubGlobal('fetch', mockFetch)
+  mockFetch.mockReset()
+  mockFetch.mockImplementation(async () => Response.json({ agentToolPermissionModeEnabled: true }))
+})
+afterEach(() => vi.unstubAllGlobals())
+
 const transaction = {
   select: () => ({ from: () => ({ where: mockSelectWhere }) }),
   update: () => ({ set: mockSet }),
+  delete: vi.fn(),
+  insert: vi.fn(),
 }
 
 describe('search replacement persistence', () => {
@@ -70,6 +91,7 @@ describe('search replacement persistence', () => {
     mockSelectWhere.mockResolvedValue([
       {
         id: 'agent-1',
+        type: 'agent',
         locked: false,
         data: {},
         subBlocks: { tools: { id: 'tools', type: 'tool-input', value: stored } },
@@ -118,3 +140,72 @@ describe('search replacement persistence', () => {
     expect(mockSet).toHaveBeenCalledTimes(1)
   })
 })
+
+describe('variable permissions at every realtime transaction write boundary', () => {
+  const block = {
+    id: 'agent-1',
+    type: 'agent',
+    name: 'Agent',
+    position: { x: 0, y: 0 },
+    locked: false,
+    subBlocks: {
+      tools: { id: 'tools', type: 'tool-input', value: [{ usageControlExpression: 'force' }] },
+    },
+    data: {},
+  }
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockTransaction.mockImplementation(
+      async (callback: (tx: typeof transaction) => Promise<void>) => callback(transaction)
+    )
+    mockSet.mockReturnValue({ where: vi.fn().mockResolvedValue(undefined) })
+    mockSelectWhere.mockImplementation(() =>
+      Object.assign(Promise.resolve([block]), { limit: async () => [block] })
+    )
+    mockFetch.mockImplementation(async () =>
+      Response.json({ agentToolPermissionModeEnabled: false })
+    )
+  })
+
+  it.each([
+    {
+      operation: WORKFLOW_OPERATIONS.REPLACE_STATE,
+      target: OPERATION_TARGETS.WORKFLOW,
+      payload: { state: { blocks: { 'agent-1': block } } },
+    },
+    {
+      operation: BLOCK_OPERATIONS.UPDATE_CANONICAL_MODE,
+      target: OPERATION_TARGETS.BLOCK,
+      payload: { id: 'agent-1', canonicalId: '0:agentToolUsageControl', canonicalMode: 'advanced' },
+    },
+    {
+      operation: BLOCK_OPERATIONS.REPLACE_CANONICAL_MODES,
+      target: OPERATION_TARGETS.BLOCK,
+      payload: {
+        id: 'agent-1',
+        data: { canonicalModes: { '0:agentToolUsageControl': 'advanced' } },
+      },
+    },
+    {
+      operation: BLOCKS_OPERATIONS.BATCH_ADD_BLOCKS,
+      target: OPERATION_TARGETS.BLOCKS,
+      payload: { blocks: [block] },
+    },
+    {
+      operation: SUBBLOCK_OPERATIONS.BATCH_UPDATE,
+      target: OPERATION_TARGETS.SUBBLOCK,
+      payload: {
+        updates: [{ blockId: 'agent-1', subblockId: 'tools', value: block.subBlocks.tools.value }],
+      },
+    },
+  ])('refuses $operation before changing any block rows', async (operation) => {
+    vi.mocked(mergeSubBlockValues).mockReturnValue(block.subBlocks)
+    await expect(
+      persistWorkflowOperation('workflow-1', { ...operation, timestamp: Date.now() })
+    ).rejects.toThrow('disabled')
+    expect(mockSet).toHaveBeenCalledTimes(1)
+    expect(transaction.delete).not.toHaveBeenCalled()
+    expect(transaction.insert).not.toHaveBeenCalled()
+  })
+})

apps/realtime/src/database/operations.ts

@@ -40,6 +40,7 @@ import {
 import { and, eq, inArray, isNull, or, sql } from 'drizzle-orm'
 import { drizzle } from 'drizzle-orm/postgres-js'
 import postgres from 'postgres'
+import { assertAgentToolPermissionModeEnabled } from '@/database/workflow-authoring'
 import { env } from '@/env'
 
 const logger = createLogger('SocketDatabase')
@@ -781,7 +782,11 @@ async function handleBlockOperationTx(
       }
 
       const existingBlock = await tx
-        .select({ data: workflowBlocks.data })
+        .select({
+          type: workflowBlocks.type,
+          subBlocks: workflowBlocks.subBlocks,
+          data: workflowBlocks.data,
+        })
         .from(workflowBlocks)
         .where(and(eq(workflowBlocks.id, payload.id), eq(workflowBlocks.workflowId, workflowId)))
         .limit(1)
@@ -793,6 +798,15 @@ async function handleBlockOperationTx(
         [payload.canonicalId]: payload.canonicalMode,
       }
 
+      if (existingBlock[0]) {
+        await assertAgentToolPermissionModeEnabled([
+          {
+            ...existingBlock[0],
+            data: { ...currentData, canonicalModes },
+          },
+        ])
+      }
+
       const updateResult = await tx
         .update(workflowBlocks)
         .set({
@@ -821,13 +835,26 @@ async function handleBlockOperationTx(
       }
 
       const existingBlock = await tx
-        .select({ data: workflowBlocks.data })
+        .select({
+          type: workflowBlocks.type,
+          subBlocks: workflowBlocks.subBlocks,
+          data: workflowBlocks.data,
+        })
         .from(workflowBlocks)
         .where(and(eq(workflowBlocks.id, payload.id), eq(workflowBlocks.workflowId, workflowId)))
         .limit(1)
 
       const currentData = (existingBlock?.[0]?.data as Record<string, unknown>) || {}
 
+      if (existingBlock[0]) {
+        await assertAgentToolPermissionModeEnabled([
+          {
+            ...existingBlock[0],
+            data: { ...currentData, canonicalModes: payload.data.canonicalModes },
+          },
+        ])
+      }
+
       const updateResult = await tx
         .update(workflowBlocks)
         .set({
@@ -975,6 +1002,8 @@ async function handleBlocksOperationTx(
           }
         })
 
+        await assertAgentToolPermissionModeEnabled(blockValues)
+
         await tx
           .insert(workflowBlocks)
           .values(blockValues)
@@ -2006,6 +2035,7 @@ async function handleSubblockOperationTx(
       const allBlocks = await tx
         .select({
           id: workflowBlocks.id,
+          type: workflowBlocks.type,
           subBlocks: workflowBlocks.subBlocks,
           locked: workflowBlocks.locked,
           data: workflowBlocks.data,
@@ -2045,6 +2075,10 @@ async function handleSubblockOperationTx(
           ? { ...currentSubBlock, value }
           : { id: subblockId, type: 'unknown', value }
 
+        if (subblockId === 'tools') {
+          await assertAgentToolPermissionModeEnabled([{ ...block, subBlocks }])
+        }
+
         await tx
           .update(workflowBlocks)
           .set({
@@ -2156,6 +2190,7 @@ async function handleWorkflowOperationTx(
       }
 
       const { blocks, edges, loops, parallels } = payload.state
+      await assertAgentToolPermissionModeEnabled(Object.values(blocks || {}))
 
       logger.info(`Replacing workflow state for ${workflowId}`, {
         blockCount: Object.keys(blocks || {}).length,

apps/realtime/src/database/workflow-authoring.test.ts

@@ -0,0 +1,90 @@
+/** @vitest-environment node */
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { assertAgentToolPermissionModeEnabled } from '@/database/workflow-authoring'
+
+const mockFetch = vi.fn()
+const variableBlock = {
+  type: 'agent',
+  subBlocks: { tools: { value: [{ usageControlExpression: 'none' }] } },
+}
+
+describe('realtime workflow authoring policy', () => {
+  beforeEach(() => {
+    vi.stubGlobal('fetch', mockFetch)
+    mockFetch.mockReset()
+  })
+  afterEach(() => vi.unstubAllGlobals())
+
+  it('does not depend on the app for fixed, empty, or non-agent edits', async () => {
+    await assertAgentToolPermissionModeEnabled([
+      { type: 'agent', subBlocks: { tools: { value: [{ usageControl: 'force' }] } } },
+      { type: 'agent' },
+      { ...variableBlock, type: 'function' },
+    ])
+    expect(mockFetch).not.toHaveBeenCalled()
+  })
+
+  it('accepts enabled authoring and checks the current policy again on the next write', async () => {
+    mockFetch
+      .mockResolvedValueOnce(Response.json({ agentToolPermissionModeEnabled: true }))
+      .mockResolvedValueOnce(Response.json({ agentToolPermissionModeEnabled: false }))
+    await expect(assertAgentToolPermissionModeEnabled([variableBlock])).resolves.toBeUndefined()
+    await expect(assertAgentToolPermissionModeEnabled([variableBlock])).rejects.toThrow('disabled')
+    expect(mockFetch).toHaveBeenCalledTimes(2)
+    expect(mockFetch).toHaveBeenCalledWith(
+      expect.stringContaining('/api/internal/workflow-authoring-policy'),
+      expect.objectContaining({
+        cache: 'no-store',
+        redirect: 'error',
+        signal: expect.any(AbortSignal),
+      })
+    )
+  })
+
+  it.each([
+    variableBlock,
+    {
+      type: 'agent',
+      subBlocks: { tools: { value: [{}] } },
+      data: { canonicalModes: { '0:agentToolUsageControl': 'advanced' } },
+    },
+    {
+      type: 'agent',
+      subBlocks: { tools: { value: [{ usageControlExpression: '' }] } },
+      data: { canonicalModes: { '0:agentToolUsageControl': 'basic' } },
+    },
+  ])('rejects active, missing, and dormant expressions when disabled: %j', async (block) => {
+    mockFetch.mockResolvedValue(Response.json({ agentToolPermissionModeEnabled: false }))
+    await expect(assertAgentToolPermissionModeEnabled([block])).rejects.toThrow('disabled')
+  })
+
+  it.each([{}, { agentToolPermissionModeEnabled: 'true' }, null])(
+    'rejects malformed policy %j',
+    async (body) => {
+      mockFetch.mockResolvedValue(Response.json(body))
+      await expect(assertAgentToolPermissionModeEnabled([variableBlock])).rejects.toThrow()
+    }
+  )
+
+  it.each([401, 403, 500, 503])('rejects HTTP %i without persisting', async (status) => {
+    mockFetch.mockResolvedValue(new Response('', { status }))
+    await expect(assertAgentToolPermissionModeEnabled([variableBlock])).rejects.toThrow(
+      'Unable to verify'
+    )
+  })
+
+  it('propagates a dropped connection and accepts a later healthy request', async () => {
+    mockFetch
+      .mockRejectedValueOnce(new TypeError('connection reset'))
+      .mockResolvedValueOnce(Response.json({ agentToolPermissionModeEnabled: true }))
+    await expect(assertAgentToolPermissionModeEnabled([variableBlock])).rejects.toThrow(
+      'connection reset'
+    )
+    await expect(assertAgentToolPermissionModeEnabled([variableBlock])).resolves.toBeUndefined()
+  })
+
+  it('rejects a truncated body', async () => {
+    mockFetch.mockResolvedValue(new Response('{"agentToolPermissionModeEnabled":'))
+    await expect(assertAgentToolPermissionModeEnabled([variableBlock])).rejects.toThrow()
+  })
+})

apps/realtime/src/database/workflow-authoring.ts

@@ -0,0 +1,33 @@
+import { workflowAuthoringPolicySchema } from '@sim/realtime-protocol/workflow-authoring'
+import {
+  type AgentToolPermissionBlock,
+  hasVariableAgentToolPermissions,
+} from '@sim/workflow-types/agent-tool-permissions'
+import { env, getBaseUrl } from '@/env'
+
+export class AgentToolPermissionModeDisabledError extends Error {
+  constructor() {
+    super('Variable agent tool permission modes are disabled')
+    this.name = 'AgentToolPermissionModeDisabledError'
+  }
+}
+
+/**
+ * The app owns flag evaluation. Only variable-bearing writes cross this service boundary;
+ * ordinary workflow edits remain local. A failed lookup aborts the pending write.
+ */
+export async function assertAgentToolPermissionModeEnabled(
+  blocks: Iterable<AgentToolPermissionBlock>
+): Promise<void> {
+  if (!hasVariableAgentToolPermissions(blocks)) return
+
+  const response = await fetch(`${getBaseUrl()}/api/internal/workflow-authoring-policy`, {
+    headers: { 'x-api-key': env.INTERNAL_API_SECRET },
+    signal: AbortSignal.timeout(5_000),
+    cache: 'no-store',
+    redirect: 'error',
+  })
+  if (!response.ok) throw new Error('Unable to verify workflow authoring policy')
+  const policy = workflowAuthoringPolicySchema.parse(await response.json())
+  if (!policy.agentToolPermissionModeEnabled) throw new AgentToolPermissionModeDisabledError()
+}

apps/realtime/src/handlers/operations.ts

@@ -15,6 +15,7 @@ import { getErrorMessage } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
 import { ZodError } from 'zod'
 import { persistWorkflowOperation } from '@/database/operations'
+import { AgentToolPermissionModeDisabledError } from '@/database/workflow-authoring'
 import type { AuthenticatedSocket } from '@/middleware/auth'
 import { checkWorkflowOperationPermission } from '@/middleware/permissions'
 import { type IRoomManager, type UserSession, workflowRoom as wf } from '@/rooms'
@@ -605,7 +606,9 @@ export function setupOperationsHandlers(socket: AuthenticatedSocket, roomManager
         socket.emit('operation-failed', {
           operationId,
           error: errorMessage,
-          retryable: !(error instanceof ZodError),
+          retryable:
+            !(error instanceof ZodError) &&
+            !(error instanceof AgentToolPermissionModeDisabledError),
         })
       }