Skip to content

Commit 22c859e

Browse files
committed
fix(agent): harden permission editing and failure recovery
1 parent 661f655 commit 22c859e

17 files changed

Lines changed: 753 additions & 41 deletions

File tree

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/** @vitest-environment node */
2+
import { OPERATION_TARGETS, SUBBLOCK_OPERATIONS } from '@sim/realtime-protocol/constants'
3+
import { beforeEach, describe, expect, it, vi } from 'vitest'
4+
5+
const { mockTransaction, mockSelectWhere, mockSet } = vi.hoisted(() => ({
6+
mockTransaction: vi.fn(),
7+
mockSelectWhere: vi.fn(),
8+
mockSet: vi.fn(),
9+
}))
10+
11+
vi.mock('@sim/audit', () => ({ AuditAction: {}, AuditResourceType: {}, recordAudit: vi.fn() }))
12+
vi.mock('@sim/db', () => ({
13+
instrumentPoolClient: vi.fn(),
14+
resolveDbUrl: vi.fn(() => 'postgres://localhost/test'),
15+
workflow: { id: 'workflow.id' },
16+
workflowBlocks: { id: 'block.id', workflowId: 'block.workflowId' },
17+
workflowEdges: {},
18+
workflowSubflows: {},
19+
}))
20+
vi.mock('@sim/db/timestamps', () => ({ withUtcTimestamps: (options: unknown) => options }))
21+
vi.mock('@sim/logger', () => ({
22+
createLogger: () => ({ info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }),
23+
}))
24+
vi.mock('@sim/platform-authz/workflow', () => ({
25+
getActiveWorkflowContext: vi.fn().mockResolvedValue({ id: 'workflow-1' }),
26+
}))
27+
vi.mock('@sim/workflow-persistence/load', () => ({
28+
loadWorkflowFromNormalizedTablesRaw: vi.fn(),
29+
}))
30+
vi.mock('@sim/workflow-persistence/subblocks', () => ({ mergeSubBlockValues: vi.fn() }))
31+
vi.mock('drizzle-orm', () => ({
32+
and: vi.fn(),
33+
eq: vi.fn(),
34+
inArray: vi.fn(),
35+
isNull: vi.fn(),
36+
or: vi.fn(),
37+
sql: vi.fn(),
38+
}))
39+
vi.mock('drizzle-orm/postgres-js', () => ({ drizzle: () => ({ transaction: mockTransaction }) }))
40+
vi.mock('postgres', () => ({ default: vi.fn() }))
41+
vi.mock('@/env', () => ({ env: { DATABASE_URL: 'postgres://localhost/test' } }))
42+
43+
import { persistWorkflowOperation } from '@/database/operations'
44+
45+
const transaction = {
46+
select: () => ({ from: () => ({ where: mockSelectWhere }) }),
47+
update: () => ({ set: mockSet }),
48+
}
49+
50+
describe('search replacement persistence', () => {
51+
const expected = [
52+
{
53+
type: 'function',
54+
params: { language: 'javascript', code: 'return 1' },
55+
usageControl: 'none',
56+
usageControlExpression: 'auto',
57+
},
58+
]
59+
const replacement = [{ ...expected[0], usageControlExpression: 'none' }]
60+
61+
beforeEach(() => {
62+
vi.clearAllMocks()
63+
mockTransaction.mockImplementation(
64+
async (callback: (tx: typeof transaction) => Promise<void>) => callback(transaction)
65+
)
66+
mockSet.mockReturnValue({ where: vi.fn().mockResolvedValue(undefined) })
67+
})
68+
69+
function replaceTools(stored: unknown, expectedValue: unknown = expected) {
70+
mockSelectWhere.mockResolvedValue([
71+
{
72+
id: 'agent-1',
73+
locked: false,
74+
data: {},
75+
subBlocks: { tools: { id: 'tools', type: 'tool-input', value: stored } },
76+
},
77+
])
78+
return persistWorkflowOperation('workflow-1', {
79+
operation: SUBBLOCK_OPERATIONS.BATCH_UPDATE,
80+
target: OPERATION_TARGETS.SUBBLOCK,
81+
timestamp: Date.now(),
82+
payload: {
83+
updates: [{ blockId: 'agent-1', subblockId: 'tools', value: replacement, expectedValue }],
84+
},
85+
})
86+
}
87+
88+
it('accepts equivalent nested tool objects after JSONB changes their key order', async () => {
89+
const stored = [
90+
{
91+
usageControlExpression: 'auto',
92+
usageControl: 'none',
93+
params: { code: 'return 1', language: 'javascript' },
94+
type: 'function',
95+
},
96+
]
97+
98+
await expect(replaceTools(stored)).resolves.toBeUndefined()
99+
expect(mockSet).toHaveBeenLastCalledWith(
100+
expect.objectContaining({
101+
subBlocks: { tools: { id: 'tools', type: 'tool-input', value: replacement } },
102+
})
103+
)
104+
})
105+
106+
it('still rejects a permission expression changed by another editor', async () => {
107+
await expect(
108+
replaceTools([{ ...expected[0], usageControlExpression: 'force' }])
109+
).rejects.toThrow('changed since replacement was planned')
110+
expect(mockSet).toHaveBeenCalledTimes(1)
111+
})
112+
113+
it('still rejects reordered tool arrays', async () => {
114+
const another = { ...expected[0], usageControlExpression: 'force' }
115+
await expect(replaceTools([another, expected[0]], [expected[0], another])).rejects.toThrow(
116+
'changed since replacement was planned'
117+
)
118+
expect(mockSet).toHaveBeenCalledTimes(1)
119+
})
120+
})

apps/realtime/src/database/operations.ts

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { isDeepStrictEqual } from 'node:util'
12
import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
23
import * as schema from '@sim/db'
34
import {
@@ -1988,10 +1989,6 @@ async function handleSubflowOperationTx(
19881989
}
19891990
}
19901991

1991-
function valuesEqual(left: unknown, right: unknown): boolean {
1992-
return JSON.stringify(left) === JSON.stringify(right)
1993-
}
1994-
19951992
// Subblock operations - targeted value updates without replacing workflow state
19961993
async function handleSubblockOperationTx(
19971994
tx: any,
@@ -2039,7 +2036,8 @@ async function handleSubblockOperationTx(
20392036
const subBlocks = { ...((block.subBlocks as Record<string, any>) || {}) }
20402037
const currentSubBlock = subBlocks[subblockId]
20412038
const currentValue = currentSubBlock?.value
2042-
if (expectedValue !== undefined && !valuesEqual(currentValue, expectedValue)) {
2039+
/** JSONB can reorder object keys; changed values and array order must still conflict. */
2040+
if (expectedValue !== undefined && !isDeepStrictEqual(currentValue, expectedValue)) {
20432041
throw new Error(`Subblock ${blockId}.${subblockId} changed since replacement was planned`)
20442042
}
20452043

apps/sim/app/workspace/[workspaceId]/home/home.tsx

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ import {
5959
searchFilterParsers,
6060
searchQueryParam,
6161
} from '@/app/workspace/[workspaceId]/home/search-params'
62+
import { useFeatureFlag } from '@/app/workspace/[workspaceId]/providers/feature-flags-provider'
6263
import { useFolders } from '@/hooks/queries/folders'
6364
import { useMarkMothershipChatRead } from '@/hooks/queries/mothership-chats'
6465
import { useWorkflows } from '@/hooks/queries/workflows'
@@ -181,6 +182,7 @@ export function Home({ chatId, userName, userId }: HomeProps) {
181182
[setSearchQueryParam, setSearchFilters]
182183
)
183184
const memberAccessAvailable = useMemberAccessAvailable()
185+
const permissionModeEnabled = useFeatureFlag('agent-tool-permission-mode')
184186
const [composerMode, setComposerMode] = useMothershipMode()
185187
const hasCheckedLandingStorageRef = useRef(false)
186188
const initialViewInputRef = useRef<HTMLDivElement>(null)
@@ -196,6 +198,7 @@ export function Home({ chatId, userName, userId }: HomeProps) {
196198
content: seed.workflowJson,
197199
filename: `${seed.workflowName}.json`,
198200
workspaceId,
201+
agentToolPermissionModeEnabled: permissionModeEnabled,
199202
nameOverride: seed.workflowName,
200203
descriptionOverride: seed.workflowDescription || undefined,
201204
createWorkflow: async ({ name, description, workspaceId }) => {
@@ -222,7 +225,7 @@ export function Home({ chatId, userName, userId }: HomeProps) {
222225
logger.error('Error creating workflow from landing workflow seed:', error)
223226
}
224227
},
225-
[workspaceId]
228+
[workspaceId, permissionModeEnabled]
226229
)
227230

228231
useEffect(() => {

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1582,7 +1582,8 @@ export const ToolInput = memo(function ToolInput({
15821582
const hasOperations =
15831583
!isCustomTool && !isMcpFamily && hasMultipleOperations(toolBlock ?? undefined)
15841584
const showToolControl = supportsToolControl && !(isMcpTool && isMcpToolUnavailable(tool))
1585-
const showCanonicalToolControl = showToolControl && permissionModeEnabled
1585+
const showCanonicalToolControl =
1586+
showToolControl && (permissionModeEnabled || toolUsageControlMode === 'advanced')
15861587
const hasToolBody =
15871588
showCanonicalToolControl || hasOperations || displaySubBlocks.length > 0
15881589

@@ -1864,7 +1865,7 @@ export const ToolInput = memo(function ToolInput({
18641865
tool={tool}
18651866
mode={toolUsageControlMode}
18661867
supportsForce={supportsForce}
1867-
disabled={disabled || isPreview}
1868+
disabled={disabled || isPreview || !permissionModeEnabled}
18681869
onFixedChange={(usageControl) =>
18691870
handleUsageControlChange(toolIndex, usageControl)
18701871
}

apps/sim/app/workspace/[workspaceId]/w/hooks/use-import-workflow.ts

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,19 @@
11
import { useCallback, useRef, useState } from 'react'
2+
import { toast } from '@sim/emcn'
23
import { createLogger } from '@sim/logger'
4+
import { getErrorMessage } from '@sim/utils/errors'
35
import { useQueryClient } from '@tanstack/react-query'
46
import { useRouter } from 'next/navigation'
57
import { usePostHog } from 'posthog-js/react'
68
import { captureEvent } from '@/lib/posthog/client'
79
import {
10+
assertWorkflowImportFeatures,
811
extractWorkflowsFromFiles,
912
extractWorkflowsFromZip,
1013
persistImportedWorkflow,
1114
sanitizePathSegment,
1215
} from '@/lib/workflows/operations/import-export'
16+
import { useFeatureFlag } from '@/app/workspace/[workspaceId]/providers/feature-flags-provider'
1317
import { useCreateFolder } from '@/hooks/queries/folders'
1418
import { folderKeys } from '@/hooks/queries/utils/folder-keys'
1519
import { invalidateWorkflowLists } from '@/hooks/queries/utils/invalidate-workflow-lists'
@@ -34,6 +38,7 @@ interface UseImportWorkflowProps {
3438
*/
3539
export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
3640
const router = useRouter()
41+
const permissionModeEnabled = useFeatureFlag('agent-tool-permission-mode')
3742
const createWorkflowMutation = useCreateWorkflow()
3843
const queryClient = useQueryClient()
3944
const createFolderMutation = useCreateFolder()
@@ -53,6 +58,7 @@ export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
5358
content,
5459
filename,
5560
workspaceId,
61+
agentToolPermissionModeEnabled: permissionModeEnabled,
5662
folderId,
5763
sortOrder,
5864
createWorkflow: async ({ name, description, workspaceId, folderId, sortOrder }) =>
@@ -68,7 +74,7 @@ export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
6874

6975
return result?.workflowId ?? null
7076
},
71-
[clearDiff, createWorkflowMutation, workspaceId]
77+
[clearDiff, createWorkflowMutation, workspaceId, permissionModeEnabled]
7278
)
7379

7480
/**
@@ -90,6 +96,9 @@ export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
9096
if (hasZip && fileArray.length === 1) {
9197
const zipFile = fileArray[0]
9298
const { workflows: extractedWorkflows, metadata } = await extractWorkflowsFromZip(zipFile)
99+
for (const workflow of extractedWorkflows) {
100+
assertWorkflowImportFeatures(workflow.content, permissionModeEnabled)
101+
}
93102

94103
const folderName = metadata?.workspaceName || zipFile.name.replace(/\.zip$/i, '')
95104
const importFolder = await createFolderMutation.mutateAsync({
@@ -191,6 +200,9 @@ export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
191200
}
192201
} else if (jsonFiles.length > 0) {
193202
const extractedWorkflows = await extractWorkflowsFromFiles(jsonFiles)
203+
for (const workflow of extractedWorkflows) {
204+
assertWorkflowImportFeatures(workflow.content, permissionModeEnabled)
205+
}
194206

195207
for (const workflow of extractedWorkflows) {
196208
try {
@@ -219,6 +231,7 @@ export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
219231
}
220232
} catch (error) {
221233
logger.error('Failed to import workflows:', error)
234+
toast.error(getErrorMessage(error, 'Failed to import workflows'))
222235
} finally {
223236
setIsImporting(false)
224237

@@ -227,7 +240,14 @@ export function useImportWorkflow({ workspaceId }: UseImportWorkflowProps) {
227240
}
228241
}
229242
},
230-
[importSingleWorkflow, workspaceId, router, createFolderMutation, queryClient]
243+
[
244+
importSingleWorkflow,
245+
workspaceId,
246+
router,
247+
createFolderMutation,
248+
queryClient,
249+
permissionModeEnabled,
250+
]
231251
)
232252

233253
return {

0 commit comments

Comments
 (0)