Skip to content

Commit d433d88

Browse files
j15zj15z
committed
fix(agent): preserve permissions across gated workflow edits
1 parent 1316299 commit d433d88

21 files changed

Lines changed: 482 additions & 64 deletions

File tree

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

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import { hasAgentToolPermissionChanges } from '@sim/workflow-types/agent-tool-permissions'
2+
import { describe, expect, it } from 'vitest'
3+
4+
const tool = {
5+
type: 'function',
6+
params: { code: 'return 1' },
7+
usageControl: 'auto',
8+
usageControlExpression: 'force',
9+
}
10+
function block(
11+
tools = [tool],
12+
modes: Record<string, string> = { '0:agentToolUsageControl': 'advanced' }
13+
) {
14+
return {
15+
id: 'agent-1',
16+
type: 'agent',
17+
subBlocks: { tools: { value: tools } },
18+
data: { canonicalModes: modes },
19+
}
20+
}
21+
22+
describe('variable permission changes', () => {
23+
it('allows retained permissions and unrelated parameter or mode edits', () => {
24+
const before = block()
25+
const after = block([{ ...tool, params: { code: 'return 2' }, usageControl: 'none' }], {
26+
...before.data.canonicalModes,
27+
model: 'advanced',
28+
})
29+
expect(hasAgentToolPermissionChanges([after], [before])).toBe(false)
30+
})
31+
32+
it('allows reordering and removing tools', () => {
33+
const second = { ...tool, usageControlExpression: 'none' }
34+
const before = block([tool, second], {
35+
'0:agentToolUsageControl': 'advanced',
36+
'1:agentToolUsageControl': 'advanced',
37+
})
38+
const after = block([second, tool], before.data.canonicalModes)
39+
expect(hasAgentToolPermissionChanges([after], [before])).toBe(false)
40+
expect(hasAgentToolPermissionChanges([block([second])], [before])).toBe(false)
41+
expect(hasAgentToolPermissionChanges([block([])], [before])).toBe(false)
42+
})
43+
44+
it('allows disabling variable mode with its expression retained', () => {
45+
expect(hasAgentToolPermissionChanges([block([tool], {})], [block()])).toBe(false)
46+
})
47+
48+
it('rejects activating a dormant expression', () => {
49+
expect(hasAgentToolPermissionChanges([block()], [block([tool], {})])).toBe(true)
50+
})
51+
52+
it('rejects changes to active or dormant expressions', () => {
53+
for (const modes of [{}, { '0:agentToolUsageControl': 'advanced' }]) {
54+
expect(
55+
hasAgentToolPermissionChanges(
56+
[block([{ ...tool, usageControlExpression: 'none' }], modes)],
57+
[block([tool], modes)]
58+
)
59+
).toBe(true)
60+
}
61+
})
62+
63+
it('rejects copying a variable tool or moving it to another block', () => {
64+
expect(hasAgentToolPermissionChanges([block([tool, tool])], [block()])).toBe(true)
65+
expect(hasAgentToolPermissionChanges([{ ...block(), id: 'new-agent' }], [block()])).toBe(true)
66+
expect(hasAgentToolPermissionChanges([block()], [])).toBe(true)
67+
})
68+
69+
it('rejects transferring a variable permission to a different integration', () => {
70+
expect(hasAgentToolPermissionChanges([block([{ ...tool, type: 'slack' }])], [block()])).toBe(
71+
true
72+
)
73+
})
74+
75+
it('does not let dormant duplicates consume the active permission match', () => {
76+
const before = block([tool, tool], { '0:agentToolUsageControl': 'advanced' })
77+
const after = block([tool, tool], { '1:agentToolUsageControl': 'advanced' })
78+
expect(hasAgentToolPermissionChanges([after], [before])).toBe(false)
79+
expect(
80+
hasAgentToolPermissionChanges(
81+
[
82+
block([tool, tool], {
83+
'0:agentToolUsageControl': 'advanced',
84+
'1:agentToolUsageControl': 'advanced',
85+
}),
86+
],
87+
[before]
88+
)
89+
).toBe(true)
90+
})
91+
})

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

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ vi.mock('drizzle-orm/postgres-js', () => ({ drizzle: () => ({ transaction: mockT
4646
vi.mock('postgres', () => ({ default: vi.fn() }))
4747
vi.mock('@/env', () => ({
4848
env: { DATABASE_URL: 'postgres://localhost/test' },
49-
getBaseUrl: () => 'http://localhost:3000',
49+
getInternalApiBaseUrl: () => 'http://localhost:3000',
5050
}))
5151

5252
import { mergeSubBlockValues } from '@sim/workflow-persistence/subblocks'
@@ -161,7 +161,14 @@ describe('variable permissions at every realtime transaction write boundary', ()
161161
)
162162
mockSet.mockReturnValue({ where: vi.fn().mockResolvedValue(undefined) })
163163
mockSelectWhere.mockImplementation(() =>
164-
Object.assign(Promise.resolve([block]), { limit: async () => [block] })
164+
Object.assign(
165+
Promise.resolve([{ ...block, subBlocks: { tools: { value: [{ type: 'function' }] } } }]),
166+
{
167+
limit: async () => [
168+
{ ...block, subBlocks: { tools: { value: [{ type: 'function' }] } } },
169+
],
170+
}
171+
)
165172
)
166173
mockFetch.mockImplementation(async () =>
167174
Response.json({ agentToolPermissionModeEnabled: false })
@@ -208,4 +215,49 @@ describe('variable permissions at every realtime transaction write boundary', ()
208215
expect(transaction.delete).not.toHaveBeenCalled()
209216
expect(transaction.insert).not.toHaveBeenCalled()
210217
})
218+
219+
it('persists a reordered tool array and its mode map in one write while the flag is off', async () => {
220+
const first = { type: 'function', usageControlExpression: 'auto' }
221+
const second = { type: 'function', usageControlExpression: 'force' }
222+
const original = {
223+
...block,
224+
subBlocks: { tools: { id: 'tools', type: 'tool-input', value: [first, second] } },
225+
data: { canonicalModes: { '1:agentToolUsageControl': 'advanced' } },
226+
}
227+
mockSelectWhere.mockResolvedValue([original])
228+
mockSet.mockReturnValue({
229+
where: () =>
230+
Object.assign(Promise.resolve(undefined), { returning: async () => [{ id: block.id }] }),
231+
})
232+
const subBlocks = { tools: { id: 'tools', type: 'tool-input', value: [second, first] } }
233+
const canonicalModes = { '0:agentToolUsageControl': 'advanced' }
234+
await expect(
235+
persistWorkflowOperation('workflow-1', {
236+
operation: BLOCK_OPERATIONS.REPLACE_CANONICAL_MODES,
237+
target: OPERATION_TARGETS.BLOCK,
238+
timestamp: Date.now(),
239+
payload: { id: block.id, subBlocks, data: { canonicalModes } },
240+
})
241+
).resolves.toBeUndefined()
242+
expect(mockFetch).not.toHaveBeenCalled()
243+
expect(mockSet).toHaveBeenLastCalledWith(
244+
expect.objectContaining({ subBlocks, data: { canonicalModes } })
245+
)
246+
})
247+
248+
it('refuses an atomic tool update inside a locked container', async () => {
249+
mockSelectWhere.mockResolvedValue([
250+
{ ...block, data: { parentId: 'container' } },
251+
{ id: 'container', type: 'loop', locked: true, data: {} },
252+
])
253+
await expect(
254+
persistWorkflowOperation('workflow-1', {
255+
operation: BLOCK_OPERATIONS.REPLACE_CANONICAL_MODES,
256+
target: OPERATION_TARGETS.BLOCK,
257+
timestamp: Date.now(),
258+
payload: { id: block.id, subBlocks: block.subBlocks, data: { canonicalModes: {} } },
259+
})
260+
).rejects.toThrow('locked')
261+
expect(mockSet).toHaveBeenCalledTimes(1)
262+
})
211263
})

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

Lines changed: 56 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -783,6 +783,7 @@ async function handleBlockOperationTx(
783783

784784
const existingBlock = await tx
785785
.select({
786+
id: workflowBlocks.id,
786787
type: workflowBlocks.type,
787788
subBlocks: workflowBlocks.subBlocks,
788789
data: workflowBlocks.data,
@@ -799,12 +800,15 @@ async function handleBlockOperationTx(
799800
}
800801

801802
if (existingBlock[0]) {
802-
await assertAgentToolPermissionModeEnabled([
803-
{
804-
...existingBlock[0],
805-
data: { ...currentData, canonicalModes },
806-
},
807-
])
803+
await assertAgentToolPermissionModeEnabled(
804+
[
805+
{
806+
...existingBlock[0],
807+
data: { ...currentData, canonicalModes },
808+
},
809+
],
810+
existingBlock
811+
)
808812
}
809813

810814
const updateResult = await tx
@@ -834,30 +838,48 @@ async function handleBlockOperationTx(
834838
throw new Error('Missing required fields for replace canonical modes operation')
835839
}
836840

837-
const existingBlock = await tx
841+
const allBlocks = await tx
838842
.select({
843+
id: workflowBlocks.id,
844+
locked: workflowBlocks.locked,
839845
type: workflowBlocks.type,
840846
subBlocks: workflowBlocks.subBlocks,
841847
data: workflowBlocks.data,
842848
})
843849
.from(workflowBlocks)
844-
.where(and(eq(workflowBlocks.id, payload.id), eq(workflowBlocks.workflowId, workflowId)))
845-
.limit(1)
850+
.where(eq(workflowBlocks.workflowId, workflowId))
851+
const blocksById = Object.fromEntries(
852+
allBlocks.map((block: { id: string; locked: boolean; data: Record<string, unknown> }) => [
853+
block.id,
854+
block,
855+
])
856+
)
857+
if (isWorkflowBlockProtected(payload.id, blocksById)) {
858+
throw new Error(`Block ${payload.id} is locked or inside a locked container`)
859+
}
860+
const existingBlock = allBlocks.filter((block: { id: string }) => block.id === payload.id)
846861

847862
const currentData = (existingBlock?.[0]?.data as Record<string, unknown>) || {}
848863

864+
const subBlocks = { ...(existingBlock[0]?.subBlocks || {}), ...(payload.subBlocks || {}) }
865+
849866
if (existingBlock[0]) {
850-
await assertAgentToolPermissionModeEnabled([
851-
{
852-
...existingBlock[0],
853-
data: { ...currentData, canonicalModes: payload.data.canonicalModes },
854-
},
855-
])
867+
await assertAgentToolPermissionModeEnabled(
868+
[
869+
{
870+
...existingBlock[0],
871+
subBlocks,
872+
data: { ...currentData, canonicalModes: payload.data.canonicalModes },
873+
},
874+
],
875+
existingBlock
876+
)
856877
}
857878

858879
const updateResult = await tx
859880
.update(workflowBlocks)
860881
.set({
882+
...(payload.subBlocks ? { subBlocks } : {}),
861883
data: {
862884
...currentData,
863885
canonicalModes: payload.data.canonicalModes,
@@ -946,7 +968,13 @@ async function handleBlocksOperationTx(
946968
if (blocks && blocks.length > 0) {
947969
// Fetch existing blocks to check for locked parents
948970
const existingBlocks = await tx
949-
.select({ id: workflowBlocks.id, locked: workflowBlocks.locked })
971+
.select({
972+
id: workflowBlocks.id,
973+
type: workflowBlocks.type,
974+
subBlocks: workflowBlocks.subBlocks,
975+
data: workflowBlocks.data,
976+
locked: workflowBlocks.locked,
977+
})
950978
.from(workflowBlocks)
951979
.where(eq(workflowBlocks.workflowId, workflowId))
952980

@@ -1002,7 +1030,7 @@ async function handleBlocksOperationTx(
10021030
}
10031031
})
10041032

1005-
await assertAgentToolPermissionModeEnabled(blockValues)
1033+
await assertAgentToolPermissionModeEnabled(blockValues, existingBlocks)
10061034

10071035
await tx
10081036
.insert(workflowBlocks)
@@ -2076,7 +2104,7 @@ async function handleSubblockOperationTx(
20762104
: { id: subblockId, type: 'unknown', value }
20772105

20782106
if (subblockId === 'tools') {
2079-
await assertAgentToolPermissionModeEnabled([{ ...block, subBlocks }])
2107+
await assertAgentToolPermissionModeEnabled([{ ...block, subBlocks }], [block])
20802108
}
20812109

20822110
await tx
@@ -2190,7 +2218,16 @@ async function handleWorkflowOperationTx(
21902218
}
21912219

21922220
const { blocks, edges, loops, parallels } = payload.state
2193-
await assertAgentToolPermissionModeEnabled(Object.values(blocks || {}))
2221+
const previousBlocks = await tx
2222+
.select({
2223+
id: workflowBlocks.id,
2224+
type: workflowBlocks.type,
2225+
subBlocks: workflowBlocks.subBlocks,
2226+
data: workflowBlocks.data,
2227+
})
2228+
.from(workflowBlocks)
2229+
.where(eq(workflowBlocks.workflowId, workflowId))
2230+
await assertAgentToolPermissionModeEnabled(Object.values(blocks || {}), previousBlocks)
21942231

21952232
logger.info(`Replacing workflow state for ${workflowId}`, {
21962233
blockCount: Object.keys(blocks || {}).length,

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,3 +88,15 @@ describe('realtime workflow authoring policy', () => {
8888
await expect(assertAgentToolPermissionModeEnabled([variableBlock])).rejects.toThrow()
8989
})
9090
})
91+
92+
it('does not consult the policy service for preserved variable permissions', async () => {
93+
const fetch = vi.fn()
94+
vi.stubGlobal('fetch', fetch)
95+
try {
96+
const before = { ...variableBlock, id: 'agent-1' }
97+
await assertAgentToolPermissionModeEnabled([before], [structuredClone(before)])
98+
expect(fetch).not.toHaveBeenCalled()
99+
} finally {
100+
vi.unstubAllGlobals()
101+
}
102+
})

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

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
import { workflowAuthoringPolicySchema } from '@sim/realtime-protocol/workflow-authoring'
22
import {
33
type AgentToolPermissionBlock,
4-
hasVariableAgentToolPermissions,
4+
hasAgentToolPermissionChanges,
55
} from '@sim/workflow-types/agent-tool-permissions'
6-
import { env, getBaseUrl } from '@/env'
6+
import { env, getInternalApiBaseUrl } from '@/env'
77

88
export class AgentToolPermissionModeDisabledError extends Error {
99
constructor() {
@@ -13,20 +13,24 @@ export class AgentToolPermissionModeDisabledError extends Error {
1313
}
1414

1515
/**
16-
* The app owns flag evaluation. Only variable-bearing writes cross this service boundary;
16+
* The app owns flag evaluation. Only new or changed variable permissions cross this boundary;
1717
* ordinary workflow edits remain local. A failed lookup aborts the pending write.
1818
*/
1919
export async function assertAgentToolPermissionModeEnabled(
20-
blocks: Iterable<AgentToolPermissionBlock>
20+
blocks: Iterable<AgentToolPermissionBlock>,
21+
previousBlocks: Iterable<AgentToolPermissionBlock> = []
2122
): Promise<void> {
22-
if (!hasVariableAgentToolPermissions(blocks)) return
23+
if (!hasAgentToolPermissionChanges(blocks, previousBlocks)) return
2324

24-
const response = await fetch(`${getBaseUrl()}/api/internal/workflow-authoring-policy`, {
25-
headers: { 'x-api-key': env.INTERNAL_API_SECRET },
26-
signal: AbortSignal.timeout(5_000),
27-
cache: 'no-store',
28-
redirect: 'error',
29-
})
25+
const response = await fetch(
26+
`${getInternalApiBaseUrl()}/api/internal/workflow-authoring-policy`,
27+
{
28+
headers: { 'x-api-key': env.INTERNAL_API_SECRET },
29+
signal: AbortSignal.timeout(5_000),
30+
cache: 'no-store',
31+
redirect: 'error',
32+
}
33+
)
3034
if (!response.ok) throw new Error('Unable to verify workflow authoring policy')
3135
const policy = workflowAuthoringPolicySchema.parse(await response.json())
3236
if (!policy.agentToolPermissionModeEnabled) throw new AgentToolPermissionModeDisabledError()

‎apps/realtime/src/env.test.ts‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
/** @vitest-environment node */
2+
import { afterEach, describe, expect, it } from 'vitest'
3+
import { env, getInternalApiBaseUrl } from '@/env'
4+
5+
const original = env.INTERNAL_API_BASE_URL
6+
7+
afterEach(() => {
8+
env.INTERNAL_API_BASE_URL = original
9+
})
10+
11+
describe('internal app service address', () => {
12+
it('uses the container service instead of the browser origin', () => {
13+
env.INTERNAL_API_BASE_URL = 'http://simstudio:3000/'
14+
expect(getInternalApiBaseUrl()).toBe('http://simstudio:3000')
15+
})
16+
17+
it('falls back to the public app origin for local and existing deployments', () => {
18+
env.INTERNAL_API_BASE_URL = undefined
19+
expect(getInternalApiBaseUrl()).toBe(env.NEXT_PUBLIC_APP_URL.replace(/\/+$/, ''))
20+
})
21+
})

0 commit comments

Comments
 (0)