Skip to content

Commit 266fb26

Browse files
Sg312Sg312
committed
fix(mothership): allow authorized organization code secrets in Build and Plan
1 parent 5d5ab7d commit 266fb26

10 files changed

Lines changed: 323 additions & 110 deletions

File tree

‎apps/sim/lib/function-execution/application/execute-chat-function.test.ts‎

Lines changed: 83 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
/** @vitest-environment node */
22
import { beforeEach, expect, it, vi } from 'vitest'
3+
import { FUNCTION_EXECUTION_DELEGATION_AUDIENCE } from '@/lib/function-execution/application/authorization'
34
import { createTrustedOrganizationCopilotPrincipal } from '@/lib/mothership/auth/application-delegation'
4-
import { FUNCTION_EXECUTION_DELEGATION_AUDIENCE } from './authorization'
5+
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
56

67
const mocks = vi.hoisted(() => ({ context: vi.fn(), authorize: vi.fn(), execute: vi.fn() }))
78
vi.mock('@/lib/mothership/chat/application/context', () => ({
@@ -14,7 +15,7 @@ vi.mock('@/lib/function-execution/execute-request', () => ({
1415
executeFunctionRequest: mocks.execute,
1516
}))
1617

17-
import { executeChatFunction } from './execute-chat-function'
18+
import { executeChatFunction } from '@/lib/function-execution/application/execute-chat-function'
1819

1920
const principal = createTrustedOrganizationCopilotPrincipal(
2021
{ userId: 'actor', organizationId: 'org', chatId: 'chat', delegationId: 'test' },
@@ -40,24 +41,91 @@ beforeEach(() => {
4041
})
4142
mocks.execute.mockResolvedValue(Response.json({ success: true }))
4243
})
43-
it('runs through the existing Function executor after fresh owner authorization without a workspace', async () => {
44-
await executeChatFunction.execute({ principal, input })
45-
expect(mocks.authorize).toHaveBeenCalledWith(
44+
it.each(['agent', 'plan'])(
45+
'runs through the existing Function executor after fresh %s owner authorization without a workspace',
46+
async (mode) => {
47+
mocks.context.mockResolvedValue({
48+
organizationId: 'org',
49+
chatId: 'chat',
50+
userId: 'actor',
51+
mode,
52+
})
53+
await executeChatFunction.execute({ principal, input })
54+
expect(mocks.authorize).toHaveBeenCalledWith(
55+
principal,
56+
expect.objectContaining({ capability: 'copilot.use', minimumRole: 'member' }),
57+
{ organizationId: 'org' }
58+
)
59+
expect(mocks.execute).toHaveBeenCalledWith(
60+
expect.anything(),
61+
expect.objectContaining({ code: 'return 1', secretScope: 'selected' }),
62+
expect.objectContaining({
63+
principal,
64+
attributedUserId: 'actor',
65+
sandboxProfile: 'mothership',
66+
})
67+
)
68+
expect(mocks.execute.mock.calls[0][1]).not.toHaveProperty('workspaceId')
69+
mocks.authorize.mockRejectedValueOnce(new Error('membership revoked'))
70+
await expect(executeChatFunction.execute({ principal, input })).rejects.toThrow(
71+
'membership revoked'
72+
)
73+
expect(mocks.execute).toHaveBeenCalledTimes(1)
74+
}
75+
)
76+
it.each(['agent', 'plan'])('accepts only authorized organization mounts in %s', async (mode) => {
77+
mocks.context.mockResolvedValue({ organizationId: 'org', chatId: 'chat', userId: 'actor', mode })
78+
const registry = new ResolvedSecretTraceRegistry([
79+
{ name: 'TOKEN', plaintext: 'test-token', encryptedValue: 'test-cipher' },
80+
])
81+
await executeChatFunction.execute({
4682
principal,
47-
expect.objectContaining({ capability: 'copilot.use', minimumRole: 'member' }),
48-
{ organizationId: 'org' }
49-
)
83+
input: {
84+
...input,
85+
resolvedSecretTraceRegistry: registry,
86+
body: { ...input.body, mountedSecrets: ['TOKEN'], envVars: { TOKEN: 'test-token' } },
87+
},
88+
})
5089
expect(mocks.execute).toHaveBeenCalledWith(
5190
expect.anything(),
52-
expect.objectContaining({ code: 'return 1', secretScope: 'selected' }),
53-
expect.objectContaining({ principal, attributedUserId: 'actor', sandboxProfile: 'mothership' })
91+
expect.objectContaining({ mountedSecrets: ['TOKEN'], envVars: { TOKEN: 'test-token' } }),
92+
expect.objectContaining({ resolvedSecretTraceRegistry: registry })
5493
)
55-
expect(mocks.execute.mock.calls[0][1]).not.toHaveProperty('workspaceId')
56-
mocks.authorize.mockRejectedValueOnce(new Error('membership revoked'))
57-
await expect(executeChatFunction.execute({ principal, input })).rejects.toThrow(
58-
'membership revoked'
94+
expect(registry.getActiveMatches()).toEqual(
95+
expect.arrayContaining([expect.objectContaining({ plaintext: 'test-token' })])
5996
)
60-
expect(mocks.execute).toHaveBeenCalledTimes(1)
97+
})
98+
it.each([
99+
{ envVars: { TOKEN: 'forged' }, mountedSecrets: ['TOKEN'] },
100+
{ envVars: { OTHER: 'test-token' }, mountedSecrets: ['OTHER'] },
101+
{ envVars: { TOKEN: 'test-token', EXTRA: 'unlisted' }, mountedSecrets: ['TOKEN'] },
102+
{ envVars: { TOKEN: 'test-token', EXTRA: 'unlisted' }, mountedSecrets: ['TOKEN', 'TOKEN'] },
103+
])('refuses values or names outside the trusted mount catalog %j', async (body) => {
104+
await expect(
105+
executeChatFunction.execute({
106+
principal,
107+
input: {
108+
...input,
109+
body: { ...input.body, ...body },
110+
resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([
111+
{ name: 'TOKEN', plaintext: 'test-token', encryptedValue: 'test-cipher' },
112+
]),
113+
},
114+
})
115+
).rejects.toThrow('authorized mount')
116+
expect(mocks.execute).not.toHaveBeenCalled()
117+
})
118+
it('refuses a mount with no trusted in-process provenance', async () => {
119+
await expect(
120+
executeChatFunction.execute({
121+
principal,
122+
input: {
123+
...input,
124+
body: { ...input.body, mountedSecrets: ['TOKEN'], envVars: { TOKEN: 'test-token' } },
125+
},
126+
})
127+
).rejects.toThrow('authorized mount')
128+
expect(mocks.execute).not.toHaveBeenCalled()
61129
})
62130
it.each([
63131
{ workspaceId: 'other' },

‎apps/sim/lib/function-execution/application/execute-chat-function.ts‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ interface ExecuteChatFunctionInput extends Omit<ExecuteFunctionInput, 'workspace
1414
chatId: string
1515
}
1616

17-
/** Organization scratch execution is owned by the current private Agent chat, without a synthetic workspace. */
17+
/** Organization scratch execution belongs to the current private Build/Plan chat. */
1818
export const executeChatFunction = defineAuthorizedChatUseCase({
1919
operation: defineWorkspaceOperation({
2020
id: 'function-executions.execute_chat',
@@ -43,9 +43,9 @@ export const executeChatFunction = defineAuthorizedChatUseCase({
4343
if (
4444
context.organizationId !== input.organizationId ||
4545
context.workspaceId ||
46-
context.mode !== 'agent'
46+
(context.mode !== 'agent' && context.mode !== 'plan')
4747
)
48-
throw new OrchestrationError('not_found', 'Organization Agent chat not found')
48+
throw new OrchestrationError('not_found', 'Organization Build or Plan chat not found')
4949
return context
5050
},
5151
authorizationOptions: {
@@ -60,17 +60,30 @@ export const executeChatFunction = defineAuthorizedChatUseCase({
6060
body.workspaceId ||
6161
body.workflowId ||
6262
body.sandboxId ||
63-
Object.keys(body.envVars).length ||
6463
body.secretScope !== 'selected' ||
65-
body.mountedSecrets?.length ||
6664
body.fileKeys?.length ||
6765
body.largeValueKeys?.length ||
6866
body.largeValueExecutionIds?.length ||
6967
body.allowLargeValueWorkflowScope
7068
)
7169
throw new OrchestrationError(
7270
'validation',
73-
'Workspace files, secrets, saved sandboxes and workflow values require an explicit workspace target'
71+
'Workspace files, saved sandboxes and workflow values require an explicit workspace target'
72+
)
73+
/** Only exact values from the server's authorized mount catalog may enter org scratch code. */
74+
const mountedNames = body.mountedSecrets ?? []
75+
if (
76+
new Set(mountedNames).size !== mountedNames.length ||
77+
Object.keys(body.envVars).length !== mountedNames.length ||
78+
mountedNames.some(
79+
(name) =>
80+
!Object.hasOwn(body.envVars, name) ||
81+
!input.resolvedSecretTraceRegistry?.recordResolved(name, body.envVars[name])
82+
)
83+
)
84+
throw new OrchestrationError(
85+
'forbidden',
86+
'Organization code secrets require an authorized mount'
7487
)
7588
const { executeFunctionRequest } = await import('@/lib/function-execution/execute-request')
7689
return executeFunctionRequest(

‎apps/sim/lib/internal/function/execute.test.ts‎

Lines changed: 34 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -107,46 +107,49 @@ describe('executeFunctionTool', () => {
107107
}),
108108
})
109109
})
110-
it('binds workspace-free scratch to the trusted organization Agent chat and strips forged owners', async () => {
111-
await executeFunctionTool({
112-
body: { code: 'return 1', workspaceId: 'forged', userId: 'forged', timeout: 1000 },
113-
headers: new Headers(),
114-
requestId: 'request',
115-
sandboxProfile: 'mothership',
116-
context: {
117-
userId: 'actor',
118-
organizationId: 'org',
119-
chatId: 'chat',
120-
requestMode: 'agent',
121-
copilotToolExecution: true,
122-
},
123-
})
124-
expect(mocks.executeChat).toHaveBeenCalledWith(
125-
expect.objectContaining({
126-
principal: expect.objectContaining({
127-
kind: 'organization_delegated',
128-
subjectUserId: 'actor',
129-
organizationId: 'org',
130-
resourceScope: { chatId: 'chat' },
131-
}),
132-
input: expect.objectContaining({
110+
it.each(['agent', 'plan'] as const)(
111+
'binds workspace-free scratch to the trusted organization %s chat and strips forged owners',
112+
async (requestMode) => {
113+
await executeFunctionTool({
114+
body: { code: 'return 1', workspaceId: 'forged', userId: 'forged', timeout: 1000 },
115+
headers: new Headers(),
116+
requestId: 'request',
117+
sandboxProfile: 'mothership',
118+
context: {
119+
userId: 'actor',
133120
organizationId: 'org',
134121
chatId: 'chat',
135-
body: expect.objectContaining({ workspaceId: undefined, userId: undefined }),
136-
}),
122+
requestMode,
123+
copilotToolExecution: true,
124+
},
137125
})
138-
)
139-
expect(mocks.execute).not.toHaveBeenCalled()
140-
expect(mocks.createPrincipal).not.toHaveBeenCalled()
141-
})
126+
expect(mocks.executeChat).toHaveBeenCalledWith(
127+
expect.objectContaining({
128+
principal: expect.objectContaining({
129+
kind: 'organization_delegated',
130+
subjectUserId: 'actor',
131+
organizationId: 'org',
132+
resourceScope: { chatId: 'chat' },
133+
}),
134+
input: expect.objectContaining({
135+
organizationId: 'org',
136+
chatId: 'chat',
137+
body: expect.objectContaining({ workspaceId: undefined, userId: undefined }),
138+
}),
139+
})
140+
)
141+
expect(mocks.execute).not.toHaveBeenCalled()
142+
expect(mocks.createPrincipal).not.toHaveBeenCalled()
143+
}
144+
)
142145
it.each([
143146
{ requestMode: 'assistant' },
144147
{ copilotToolExecution: false },
145148
{ chatId: undefined },
146149
{ organizationId: undefined },
147150
{ userId: undefined },
148151
{ workflowId: 'workflow' },
149-
])('refuses untrusted or non-Agent organization scope %j', async (override) => {
152+
])('refuses untrusted or non-Build/Plan organization scope %j', async (override) => {
150153
await expect(
151154
executeFunctionTool({
152155
body: { code: 'return 1' },
@@ -162,7 +165,7 @@ describe('executeFunctionTool', () => {
162165
...override,
163166
},
164167
})
165-
).rejects.toThrow('trusted Agent chat scope')
168+
).rejects.toThrow('trusted Build or Plan chat scope')
166169
expect(mocks.executeChat).not.toHaveBeenCalled()
167170
expect(mocks.execute).not.toHaveBeenCalled()
168171
})

‎apps/sim/lib/internal/function/execute.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,14 +45,14 @@ export async function executeFunctionTool(input: ExecuteFunctionToolInput): Prom
4545
if (!context.workspaceId) {
4646
if (
4747
context.copilotToolExecution !== true ||
48-
context.requestMode !== 'agent' ||
48+
(context.requestMode !== 'agent' && context.requestMode !== 'plan') ||
4949
!context.organizationId ||
5050
!context.chatId ||
5151
!context.userId ||
5252
context.workflowId ||
5353
sandboxProfile !== 'mothership'
5454
)
55-
throw new Error('Organization Function execution requires trusted Agent chat scope')
55+
throw new Error('Organization Function execution requires trusted Build or Plan chat scope')
5656
const principal = createTrustedOrganizationCopilotPrincipal(
5757
{
5858
userId: context.userId,

‎apps/sim/lib/mothership/generated/integration-catalog.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ export const IntegrationCatalogContext = z.object({
3535
export type IntegrationCatalogContext = z.infer<typeof IntegrationCatalogContext>;
3636

3737
export const IntegrationCatalogRequest = IntegrationCatalogContext.extend({
38-
mode: z.enum(["agent", "assistant"]),
38+
mode: z.enum(["agent", "assistant", "plan"]),
3939
workspaceId: z.uuid().optional(),
4040
query: z.string().max(2_000).optional(),
4141
service: z.string().max(200).optional(),

‎apps/sim/lib/mothership/integrations/application/catalog.test.ts‎

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
createTrustedCopilotPrincipal,
77
createTrustedOrganizationCopilotPrincipal,
88
} from '@/lib/mothership/auth/application-delegation'
9+
import { IntegrationCatalogRequest } from '@/lib/mothership/generated/integration-catalog'
910
import {
1011
INTEGRATION_CATALOG_AUDIENCE,
1112
projectIntegrationCatalog,
@@ -202,19 +203,27 @@ describe('integration catalog projection', () => {
202203
})
203204
})
204205
describe('catalog authorization', () => {
205-
it('preserves integration discovery for agent conversations', async () => {
206-
queueChat('agent')
207-
const result = await readIntegrationCatalog.execute({
208-
principal: principal(),
209-
input: { ...input, mode: 'agent', service: 'google-email', query: 'email' },
210-
})
211-
expect(result.operations.map((operation) => operation.toolId)).toEqual(['gmail_send'])
212-
expect(mocks.build).toHaveBeenCalledWith(
213-
'actor',
214-
{ schemaSurface: 'copilot', organizationId: 'org-1' },
215-
undefined
216-
)
217-
})
206+
it.each(['agent', 'plan'] as const)(
207+
'preserves integration discovery for %s conversations',
208+
async (mode) => {
209+
queueChat(mode)
210+
const result = await readIntegrationCatalog.execute({
211+
principal: principal(),
212+
input: IntegrationCatalogRequest.parse({
213+
...input,
214+
mode,
215+
service: 'google-email',
216+
query: 'email',
217+
}),
218+
})
219+
expect(result.operations.map((operation) => operation.toolId)).toEqual(['gmail_send'])
220+
expect(mocks.build).toHaveBeenCalledWith(
221+
'actor',
222+
{ schemaSurface: 'copilot', organizationId: 'org-1' },
223+
undefined
224+
)
225+
}
226+
)
218227

219228
it('rejects Search Assistant discovery before building native or MCP catalogs', async () => {
220229
queueChat()

0 commit comments

Comments
 (0)