fix(agent): keep tool modes on their tools across API and concurrent edits

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

@@ -1,5 +1,10 @@
 /** @vitest-environment node */
-import { OPERATION_TARGETS, SUBBLOCK_OPERATIONS } from '@sim/realtime-protocol/constants'
+import {
+  BLOCK_OPERATIONS,
+  OPERATION_TARGETS,
+  SUBBLOCK_OPERATIONS,
+} from '@sim/realtime-protocol/constants'
+import { getToolInputIdentity } from '@sim/realtime-protocol/tool-input'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const { mockTransaction, mockSelectWhere, mockSet } = vi.hoisted(() => ({
@@ -98,7 +103,7 @@ describe('search replacement persistence', () => {
       },
     ]
 
-    await expect(replaceTools(stored)).resolves.toBeUndefined()
+    await expect(replaceTools(stored)).resolves.toEqual({ applied: true })
     expect(mockSet).toHaveBeenLastCalledWith(
       expect.objectContaining({
         subBlocks: { tools: { id: 'tools', type: 'tool-input', value: replacement } },
@@ -153,7 +158,7 @@ describe('subblock update with canonical modes persistence', () => {
   }
 
   it('writes the subblock value and replaces canonical modes in one block update', async () => {
-    await expect(updateTools({})).resolves.toBeUndefined()
+    await expect(updateTools({})).resolves.toEqual({ applied: true })
 
     expect(mockSet).toHaveBeenCalledTimes(2)
     expect(mockSet).toHaveBeenLastCalledWith(
@@ -169,3 +174,92 @@ describe('subblock update with canonical modes persistence', () => {
     expect(mockSet).toHaveBeenCalledTimes(1)
   })
 })
+
+describe('tool-scoped canonical mode persistence', () => {
+  const jira = { type: 'jira', operation: 'read-bulk', params: { manualProjectId: 'MAN' } }
+  const wikipedia = { type: 'wikipedia', operation: 'wikipedia_search', params: { query: 'Sim' } }
+  const mockReturning = vi.fn()
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockTransaction.mockImplementation(
+      async (callback: (tx: typeof transaction) => Promise<void>) => callback(transaction)
+    )
+    mockReturning.mockResolvedValue([{ id: 'agent-1' }])
+    mockSet.mockReturnValue({
+      where: vi.fn(() => Object.assign(Promise.resolve(undefined), { returning: mockReturning })),
+    })
+  })
+
+  function toggle(storedTools: unknown[], payload: Record<string, unknown>) {
+    mockSelectWhere.mockReturnValue({
+      limit: vi.fn().mockResolvedValue([
+        {
+          data: { canonicalModes: { '0:projectId': 'basic' } },
+          subBlocks: { tools: { id: 'tools', type: 'tool-input', value: storedTools } },
+        },
+      ]),
+    })
+    return persistWorkflowOperation('workflow-1', {
+      operation: BLOCK_OPERATIONS.UPDATE_CANONICAL_MODE,
+      target: OPERATION_TARGETS.BLOCK,
+      timestamp: Date.now(),
+      payload: { id: 'agent-1', canonicalMode: 'advanced', ...payload },
+    })
+  }
+
+  function toolRef(toolIndex: number, tool: unknown) {
+    return { subblockId: 'tools', toolIndex, identity: getToolInputIdentity(tool) ?? {} }
+  }
+
+  it('applies a mode while its tool still holds the toggled position', async () => {
+    await expect(
+      toggle([wikipedia, jira], {
+        canonicalId: '1:agentToolUsageControl',
+        toolRef: toolRef(1, jira),
+      })
+    ).resolves.toEqual({ applied: true })
+
+    expect(mockSet).toHaveBeenLastCalledWith(
+      expect.objectContaining({
+        data: {
+          canonicalModes: { '0:projectId': 'basic', '1:agentToolUsageControl': 'advanced' },
+        },
+      })
+    )
+  })
+
+  it('refuses a mode whose tool another editor moved first', async () => {
+    await expect(
+      toggle([jira, wikipedia], {
+        canonicalId: '1:agentToolUsageControl',
+        toolRef: toolRef(1, jira),
+      })
+    ).resolves.toEqual({ applied: false })
+
+    expect(mockSet).toHaveBeenCalledTimes(1)
+  })
+
+  it('refuses a mode whose tool another editor removed first', async () => {
+    await expect(
+      toggle([wikipedia], { canonicalId: '1:agentToolUsageControl', toolRef: toolRef(1, jira) })
+    ).resolves.toEqual({ applied: false })
+
+    expect(mockSet).toHaveBeenCalledTimes(1)
+  })
+
+  it('refuses a key that does not match the referenced position', async () => {
+    await expect(
+      toggle([wikipedia, jira], {
+        canonicalId: '0:agentToolUsageControl',
+        toolRef: toolRef(1, jira),
+      })
+    ).resolves.toEqual({ applied: false })
+  })
+
+  it('applies a mode sent without a tool reference as before', async () => {
+    await expect(toggle([], { canonicalId: 'files' })).resolves.toEqual({ applied: true })
+
+    expect(mockSet).toHaveBeenCalledTimes(2)
+  })
+})

apps/realtime/src/database/operations.ts

@@ -23,6 +23,8 @@ import {
   VARIABLE_OPERATIONS,
   WORKFLOW_OPERATIONS,
 } from '@sim/realtime-protocol/constants'
+import { isToolInputRefCurrent } from '@sim/realtime-protocol/tool-input'
+import { isRecordLike } from '@sim/utils/object'
 import { randomFloat } from '@sim/utils/random'
 import { loadWorkflowFromNormalizedTablesRaw } from '@sim/workflow-persistence/load'
 import { mergeSubBlockValues } from '@sim/workflow-persistence/subblocks'
@@ -374,7 +376,15 @@ export async function getWorkflowState(workflowId: string) {
   }
 }
 
-export async function persistWorkflowOperation(workflowId: string, operation: any) {
+export interface PersistWorkflowOperationResult {
+  /** False when the operation was a stale no-op that must not be broadcast to other editors. */
+  applied: boolean
+}
+
+export async function persistWorkflowOperation(
+  workflowId: string,
+  operation: any
+): Promise<PersistWorkflowOperationResult> {
   const startTime = Date.now()
   try {
     const { operation: op, target, payload, timestamp, userId } = operation
@@ -392,6 +402,7 @@ export async function persistWorkflowOperation(workflowId: string, operation: an
       })
     }
 
+    let applied = true
     await db.transaction(async (tx) => {
       // This UPDATE is also this workflow's write-serialization point, not
       // just a timestamp bump: it takes a row lock on `workflow` for the
@@ -408,9 +419,11 @@ export async function persistWorkflowOperation(workflowId: string, operation: an
         .where(eq(workflow.id, workflowId))
 
       switch (target) {
-        case OPERATION_TARGETS.BLOCK:
-          await handleBlockOperationTx(tx, workflowId, op, payload)
+        case OPERATION_TARGETS.BLOCK: {
+          const result = await handleBlockOperationTx(tx, workflowId, op, payload)
+          applied = result?.applied ?? true
           break
+        }
         case OPERATION_TARGETS.BLOCKS:
           await handleBlocksOperationTx(tx, workflowId, op, payload)
           break
@@ -457,6 +470,8 @@ export async function persistWorkflowOperation(workflowId: string, operation: an
         workflowId: `${workflowId.substring(0, 8)}...`,
       })
     }
+
+    return { applied }
   } catch (error) {
     const duration = Date.now() - startTime
     logger.error(
@@ -504,12 +519,17 @@ async function auditWorkflowLockToggle(workflowId: string, actorId: string): Pro
   })
 }
 
+function getSubBlockValue(subBlocks: unknown, subblockId: string): unknown {
+  const subBlock = isRecordLike(subBlocks) ? subBlocks[subblockId] : undefined
+  return isRecordLike(subBlock) ? subBlock.value : undefined
+}
+
 async function handleBlockOperationTx(
   tx: any,
   workflowId: string,
   operation: string,
   payload: any
-) {
+): Promise<PersistWorkflowOperationResult | undefined> {
   switch (operation) {
     case BLOCK_OPERATIONS.UPDATE_POSITION: {
       if (!payload.id || !payload.position) {
@@ -782,11 +802,25 @@ async function handleBlockOperationTx(
       }
 
       const existingBlock = await tx
-        .select({ data: workflowBlocks.data })
+        .select({ data: workflowBlocks.data, subBlocks: workflowBlocks.subBlocks })
         .from(workflowBlocks)
         .where(and(eq(workflowBlocks.id, payload.id), eq(workflowBlocks.workflowId, workflowId)))
         .limit(1)
 
+      if (
+        payload.toolRef &&
+        !isToolInputRefCurrent(
+          getSubBlockValue(existingBlock?.[0]?.subBlocks, payload.toolRef.subblockId),
+          payload.canonicalId,
+          payload.toolRef
+        )
+      ) {
+        logger.debug(
+          `Skipped stale tool canonical mode: ${payload.id} -> ${payload.canonicalId}: ${payload.canonicalMode}`
+        )
+        return { applied: false }
+      }
+
       const currentData = (existingBlock?.[0]?.data as Record<string, unknown>) || {}
       const currentCanonicalModes = (currentData.canonicalModes as Record<string, unknown>) || {}
       const canonicalModes = {

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

@@ -179,3 +179,59 @@ describe('workflow operation ACL', () => {
     })
   })
 })
+
+describe('tool-scoped canonical mode broadcast', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockAssertMutable.mockResolvedValue(undefined)
+    mockAuthorizeWorkflow.mockResolvedValue({ allowed: true, workspacePermission: 'write' })
+  })
+
+  function toolModeToggle(operationId: string) {
+    return {
+      operationId,
+      operation: 'update-canonical-mode',
+      target: 'block',
+      timestamp: Date.now(),
+      payload: {
+        id: BLOCK_ID,
+        canonicalId: '1:agentToolUsageControl',
+        canonicalMode: 'advanced',
+        toolRef: { subblockId: 'tools', toolIndex: 1, identity: { type: 'jira' } },
+      },
+    }
+  }
+
+  it('confirms a stale toggle to its sender without broadcasting it', async () => {
+    mockPersist.mockResolvedValue({ applied: false })
+    const { socket, handlers, toEmit } = setup('sock-tool-mode-1', 'write')
+
+    await handlers['workflow-operation'](toolModeToggle('op-stale'))
+
+    expect(mockPersist).toHaveBeenCalledWith(
+      WORKFLOW_ID,
+      expect.objectContaining({
+        payload: expect.objectContaining({
+          toolRef: { subblockId: 'tools', toolIndex: 1, identity: { type: 'jira' } },
+        }),
+      })
+    )
+    expect(toEmit).not.toHaveBeenCalled()
+    expect(socket.emit).toHaveBeenCalledWith(
+      'operation-confirmed',
+      expect.objectContaining({ operationId: 'op-stale' })
+    )
+  })
+
+  it('broadcasts a toggle that applied', async () => {
+    mockPersist.mockResolvedValue({ applied: true })
+    const { handlers, toEmit } = setup('sock-tool-mode-2', 'write')
+
+    await handlers['workflow-operation'](toolModeToggle('op-applied'))
+
+    expect(toEmit).toHaveBeenCalledWith(
+      'workflow-operation',
+      expect.objectContaining({ operation: 'update-canonical-mode' })
+    )
+  })
+})

apps/realtime/src/handlers/operations.ts

@@ -566,7 +566,7 @@ export function setupOperationsHandlers(socket: AuthenticatedSocket, roomManager
       }
 
       // For non-position operations, persist first then broadcast
-      await persistWorkflowOperation(workflowId, {
+      const persisted = await persistWorkflowOperation(workflowId, {
         operation,
         target,
         payload,
@@ -576,21 +576,23 @@ export function setupOperationsHandlers(socket: AuthenticatedSocket, roomManager
 
       await roomManager.updateRoomLastModified(wf(workflowId))
 
-      const broadcastData = {
-        operation,
-        target,
-        payload,
-        timestamp: operationTimestamp,
-        senderId: socket.id,
-        userId: session.userId,
-        userName: session.userName,
-        metadata: {
-          workflowId,
-          operationId: generateId(),
-        },
-      }
+      if (persisted?.applied !== false) {
+        const broadcastData = {
+          operation,
+          target,
+          payload,
+          timestamp: operationTimestamp,
+          senderId: socket.id,
+          userId: session.userId,
+          userName: session.userName,
+          metadata: {
+            workflowId,
+            operationId: generateId(),
+          },
+        }
 
-      socket.to(workflowId).emit('workflow-operation', broadcastData)
+        socket.to(workflowId).emit('workflow-operation', broadcastData)
+      }
 
       if (operationId) {
         socket.emit('operation-confirmed', {