Skip to content

Commit 347e721

Browse files
Sg312Sg312
committed
fix(mothership): keep workbench policy out of the public CLI
1 parent a2eafd9 commit 347e721

21 files changed

Lines changed: 232 additions & 173 deletions

File tree

‎apps/sim/lib/execution/remote-sandbox/session-cli.test.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,36 @@ async function fixture() {
106106
}
107107

108108
describe('versioned workbench CLI installation', () => {
109+
it('installs and verifies the runtime before activating the private entrypoint', async () => {
110+
const f = await fixture()
111+
const original = f.session('paired')
112+
const runtime = {
113+
path: join(dirname(original.cli.path), 'runtime.mjs'),
114+
content: 'export const result = "public-runtime";',
115+
}
116+
const session = {
117+
...original,
118+
cli: {
119+
...original.cli,
120+
content:
121+
'#!/usr/bin/env node\nimport { result } from "./runtime.mjs"; console.log(result);\n',
122+
runtime,
123+
},
124+
}
125+
await f.ensure(session)
126+
expect((await exec('sim', [], { env: f.environment(session) })).stdout.trim()).toBe(
127+
'public-runtime'
128+
)
129+
const writes = f.writes.length
130+
await f.ensure(session)
131+
expect(f.writes).toHaveLength(writes)
132+
await writeFile(runtime.path, 'damaged-runtime')
133+
await f.ensure(session)
134+
expect((await exec('sim', [], { env: f.environment(session) })).stdout.trim()).toBe(
135+
'public-runtime'
136+
)
137+
})
138+
109139
it('runs the installed CLI and reuses complete bytes without rewriting or relinking', async () => {
110140
const f = await fixture()
111141
const session = f.session('one')

‎apps/sim/lib/execution/remote-sandbox/session-cli.ts‎

Lines changed: 35 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,20 @@ import type { SandboxHandle, SandboxSessionRequest } from '@/lib/execution/remot
55

66
/** A valid release is reused; activation repairs only its own versioned launcher. */
77
const CHECK_CLI_COMMAND = `python3 - <<'SIM_SESSION_CLI'
8-
import hashlib, os, stat, sys, tempfile
8+
import hashlib, json, os, stat, sys, tempfile
99
path = os.environ['SIM_CLI_PATH']
10-
if os.path.islink(path):
11-
sys.exit(10)
12-
try:
10+
for artifact, expected in json.loads(os.environ['SIM_CLI_ARTIFACTS']).items():
11+
if os.path.islink(artifact):
12+
sys.exit(10)
1313
digest = hashlib.sha256()
14-
with open(path, 'rb') as source:
15-
for chunk in iter(lambda: source.read(65536), b''):
16-
digest.update(chunk)
17-
except FileNotFoundError:
18-
sys.exit(10)
19-
if digest.hexdigest() != os.environ['SIM_CLI_SHA256']:
20-
sys.exit(10)
14+
try:
15+
with open(artifact, 'rb') as source:
16+
for chunk in iter(lambda: source.read(65536), b''):
17+
digest.update(chunk)
18+
except FileNotFoundError:
19+
sys.exit(10)
20+
if digest.hexdigest() != expected:
21+
sys.exit(10)
2122
if stat.S_IMODE(os.stat(path).st_mode) != 0o755:
2223
os.chmod(path, 0o755)
2324
launcher = os.path.join(os.path.dirname(path), 'sim')
@@ -53,7 +54,14 @@ export async function ensureSessionCli(
5354
const result = await sandbox.runCommand(CHECK_CLI_COMMAND, {
5455
envs: {
5556
SIM_CLI_PATH: cli.path,
56-
SIM_CLI_SHA256: createHash('sha256').update(cli.content).digest('hex'),
57+
SIM_CLI_ARTIFACTS: JSON.stringify(
58+
Object.fromEntries(
59+
[cli, ...(cli.runtime ? [cli.runtime] : [])].map(({ path, content }) => [
60+
path,
61+
createHash('sha256').update(content).digest('hex'),
62+
])
63+
)
64+
),
5765
},
5866
timeoutMs: Math.min(30_000, remainingMs()),
5967
maxOutputBytes: 64 * 1024,
@@ -70,18 +78,20 @@ export async function ensureSessionCli(
7078
)
7179
}
7280
if (await check()) return
73-
await withSandboxFilePublication(
74-
sandbox,
75-
cli.path,
76-
{
77-
overwrite: true,
78-
followSymlinks: false,
79-
executable: true,
80-
rootUser: false,
81-
signal,
82-
timeoutMs: remainingMs(),
83-
},
84-
(staged) => sandbox.writeFile(staged, cli.content)
85-
)
81+
for (const artifact of [...(cli.runtime ? [cli.runtime] : []), cli]) {
82+
await withSandboxFilePublication(
83+
sandbox,
84+
artifact.path,
85+
{
86+
overwrite: true,
87+
followSymlinks: false,
88+
executable: artifact === cli,
89+
rootUser: false,
90+
signal,
91+
timeoutMs: remainingMs(),
92+
},
93+
(staged) => sandbox.writeFile(staged, artifact.content)
94+
)
95+
}
8696
if (!(await check())) throw new Error('Workbench CLI installation could not be verified')
8797
}

‎apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,11 @@ function fakeSandbox(id: string): { handle: SandboxHandle; calls: FakeSandboxCal
9494
: createHash('sha256')
9595
.update(typeof content === 'string' ? content : new Uint8Array(content))
9696
.digest('hex')
97-
return { stdout: '', stderr: '', exitCode: hash === options.envs?.SIM_CLI_SHA256 ? 0 : 10 }
97+
return {
98+
stdout: '',
99+
stderr: '',
100+
exitCode: hash === JSON.parse(options.envs?.SIM_CLI_ARTIFACTS ?? '{}')[cliPath] ? 0 : 10,
101+
}
98102
}
99103
const stage = options.envs?.SIM_FILE_STAGE
100104
const target = options.envs?.SIM_FILE_TARGET

‎apps/sim/lib/execution/remote-sandbox/types.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ export interface SandboxSessionRequest {
8888
/** Stable identity of the session (e.g. one per Mothership chat). */
8989
key: string
9090
/** Deployment-owned CLI artifact. Its versioned directory is prepended to this execution's PATH. */
91-
cli?: { path: string; content: string }
91+
cli?: { path: string; content: string; runtime?: { path: string; content: string } }
9292
/** Extra environment variables present on every execution in the session. */
9393
envs?: Record<string, string>
9494
}

‎apps/sim/lib/function-execution/execute-request.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2270,6 +2270,7 @@ export async function executeFunctionRequest(
22702270
sessionKey: sandboxSessionKey,
22712271
workspaceId,
22722272
userId: auth.attributedUserId,
2273+
signal: executionSignal,
22732274
})
22742275
: undefined
22752276
sourceCodeForErrors = sourceCode ?? code

‎apps/sim/lib/mothership/generated/workbench.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
// GENERATED — do not edit. Source of truth: mothership worker packages/contracts/src/workbench.ts
2+
// Regenerate with `bun run contracts:sync` in the worker.
3+
4+
import { z } from "zod";
5+
6+
/** Executable bootstrap is served only on the authenticated Sim → worker connection. */
7+
export const WorkbenchBootstrap = z.strictObject({
8+
version: z.literal(1),
9+
entrypoint: z.string().min(1).max(65_536),
10+
});
11+
export type WorkbenchBootstrap = z.infer<typeof WorkbenchBootstrap>;

‎apps/sim/lib/mothership/tools/sandbox-session.test.ts‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,23 @@
11
/** @vitest-environment node */
22
import { beforeEach, describe, expect, it, vi } from 'vitest'
33

4-
const { read, mint } = vi.hoisted(() => ({ read: vi.fn(), mint: vi.fn() }))
4+
const { read, mint, fetchBootstrap, baseURL } = vi.hoisted(() => ({
5+
read: vi.fn(),
6+
mint: vi.fn(),
7+
fetchBootstrap: vi.fn(),
8+
baseURL: vi.fn(),
9+
}))
510
vi.mock('node:fs/promises', () => ({ readFile: read }))
611
vi.mock('@/lib/mothership/chat/delegation', () => ({ mintDelegationToken: mint }))
712
vi.mock('@/lib/core/config/env', () => ({
813
env: { MOTHERSHIP_SANDBOX_CLI_ENDPOINT: 'https://sim.test' },
914
}))
1015
vi.mock('@/lib/core/utils/urls', () => ({ getBaseUrl: () => 'https://unused.test' }))
16+
vi.mock('@/lib/mothership/request/go/fetch', () => ({ fetchGo: fetchBootstrap }))
17+
vi.mock('@/lib/mothership/request/headers', () => ({
18+
mothershipRequestHeaders: () => ({ 'x-api-key': 'worker-test-key' }),
19+
}))
20+
vi.mock('@/lib/mothership/server/agent-url', () => ({ getMothershipBaseURL: baseURL }))
1121

1222
import { buildMothershipSandboxSession } from '@/lib/mothership/tools/sandbox-session'
1323

@@ -17,16 +27,30 @@ describe('deployment-owned workbench tooling', () => {
1727
beforeEach(() => {
1828
vi.resetAllMocks()
1929
mint.mockResolvedValue('test-delegation')
30+
baseURL.mockResolvedValue('https://worker.test')
31+
fetchBootstrap.mockImplementation(async () =>
32+
Response.json({ version: 1, entrypoint: 'private-entry' })
33+
)
2034
})
2135

2236
it('stages the current deployment bundle and keeps credentials out of installed files', async () => {
2337
read.mockResolvedValueOnce('bundle-one').mockResolvedValueOnce('bundle-two')
2438
const first = await buildMothershipSandboxSession(request)
2539
const second = await buildMothershipSandboxSession(request)
2640
expect(first.cli).toMatchObject({
27-
content: 'bundle-one',
41+
content: 'private-entry',
2842
path: expect.stringMatching(/^\/home\/user\/\.sim-cli\/[a-f0-9]{64}\/cli\.mjs$/),
43+
runtime: { content: 'bundle-one', path: expect.stringMatching(/\/runtime\.mjs$/) },
2944
})
45+
expect(read).toHaveBeenCalledWith(expect.stringContaining('/dist/runtime.js'), 'utf8')
46+
expect(baseURL).toHaveBeenCalledWith({ userId: 'user' })
47+
expect(fetchBootstrap).toHaveBeenCalledWith(
48+
'https://worker.test/api/workbench/bootstrap',
49+
expect.objectContaining({
50+
headers: { 'x-api-key': 'worker-test-key' },
51+
redirect: 'error',
52+
})
53+
)
3054
expect(second.cli?.path).not.toBe(first.cli?.path)
3155
expect(JSON.stringify(first.cli)).not.toContain('test-delegation')
3256
expect(first.envs).toEqual({
@@ -41,4 +65,34 @@ describe('deployment-owned workbench tooling', () => {
4165
await expect(buildMothershipSandboxSession(request)).rejects.toThrow('missing workbench bundle')
4266
expect(mint).not.toHaveBeenCalled()
4367
})
68+
69+
it('pins both policy and runtime versions and refuses unavailable or malformed private bootstrap', async () => {
70+
read.mockResolvedValue('same-public-runtime')
71+
const first = await buildMothershipSandboxSession(request)
72+
fetchBootstrap.mockResolvedValueOnce(
73+
Response.json({ version: 1, entrypoint: 'private-entry-v2' })
74+
)
75+
const second = await buildMothershipSandboxSession(request)
76+
expect(second.cli?.path).not.toBe(first.cli?.path)
77+
for (const response of [
78+
new Response('Unavailable', { status: 503 }),
79+
Response.json({ version: 2, entrypoint: 'unknown' }),
80+
]) {
81+
mint.mockClear()
82+
fetchBootstrap.mockResolvedValueOnce(response)
83+
await expect(buildMothershipSandboxSession(request)).rejects.toThrow()
84+
expect(mint).not.toHaveBeenCalled()
85+
}
86+
})
87+
88+
it('cancels private bootstrap loading before provider or credential work', async () => {
89+
const controller = new AbortController()
90+
controller.abort(new Error('Stopped'))
91+
await expect(
92+
buildMothershipSandboxSession({ ...request, signal: controller.signal })
93+
).rejects.toThrow('Stopped')
94+
expect(read).not.toHaveBeenCalled()
95+
expect(fetchBootstrap).not.toHaveBeenCalled()
96+
expect(mint).not.toHaveBeenCalled()
97+
})
4498
})

‎apps/sim/lib/mothership/tools/sandbox-session.ts‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,51 @@ import { env } from '@/lib/core/config/env'
77
import { getBaseUrl } from '@/lib/core/utils/urls'
88
import type { SandboxSessionRequest } from '@/lib/execution/remote-sandbox/types'
99
import { mintDelegationToken } from '@/lib/mothership/chat/delegation'
10+
import { WorkbenchBootstrap } from '@/lib/mothership/generated/workbench'
11+
import { fetchGo } from '@/lib/mothership/request/go/fetch'
12+
import { mothershipRequestHeaders } from '@/lib/mothership/request/headers'
13+
import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url'
1014

1115
const logger = createLogger('MothershipSandboxSession')
1216

13-
/** The workbench CLI ships with Sim, so reconnect cannot silently reuse a different release. */
14-
async function workbenchCli(): Promise<{ path: string; content: string }> {
17+
/** Public runtime and private bootstrap share an immutable release directory. */
18+
async function workbenchCli(
19+
userId: string,
20+
signal?: AbortSignal
21+
): Promise<NonNullable<SandboxSessionRequest['cli']>> {
1522
const cwd = process.cwd()
1623
const path = resolve(
1724
cwd,
1825
cwd.endsWith('/apps/sim')
19-
? '../../packages/sim-cli/dist/workbench.js'
20-
: 'packages/sim-cli/dist/workbench.js'
26+
? '../../packages/sim-cli/dist/runtime.js'
27+
: 'packages/sim-cli/dist/runtime.js'
2128
)
22-
const content = await readFile(path, 'utf8')
23-
const digest = createHash('sha256').update(content).digest('hex')
24-
return { path: `/home/user/.sim-cli/${digest}/cli.mjs`, content }
29+
signal?.throwIfAborted()
30+
const runtime = await readFile(path, 'utf8')
31+
const baseURL = await getMothershipBaseURL({ userId })
32+
const deadline = AbortSignal.timeout(15_000)
33+
const response = await fetchGo(`${baseURL}/api/workbench/bootstrap`, {
34+
headers: mothershipRequestHeaders(),
35+
redirect: 'error',
36+
signal: signal ? AbortSignal.any([signal, deadline]) : deadline,
37+
spanName: 'sim → worker /api/workbench/bootstrap',
38+
operation: 'workbench_bootstrap',
39+
})
40+
if (!response.ok) {
41+
await response.body?.cancel()
42+
throw new Error('Mothership workbench bootstrap is unavailable')
43+
}
44+
const { entrypoint } = WorkbenchBootstrap.parse(await response.json())
45+
signal?.throwIfAborted()
46+
const digest = createHash('sha256')
47+
.update(JSON.stringify([runtime, entrypoint]))
48+
.digest('hex')
49+
const directory = `/home/user/.sim-cli/${digest}`
50+
return {
51+
path: `${directory}/cli.mjs`,
52+
content: entrypoint,
53+
runtime: { path: `${directory}/runtime.mjs`, content: runtime },
54+
}
2555
}
2656

2757
/**
@@ -40,8 +70,9 @@ export async function buildMothershipSandboxSession(args: {
4070
sessionKey: string
4171
workspaceId: string
4272
userId: string
73+
signal?: AbortSignal
4374
}): Promise<SandboxSessionRequest> {
44-
const cli = await workbenchCli()
75+
const cli = await workbenchCli(args.userId, args.signal)
4576
let cliEnvs: Record<string, string> | undefined
4677
try {
4778
const apiKey = await mintDelegationToken({

‎apps/sim/next.config.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ const nextConfig: NextConfig = {
142142
'/*': [
143143
'./lib/execution/sandbox/bundles/*.cjs',
144144
'./node_modules/ws/**/*',
145-
'../../packages/sim-cli/dist/workbench.js',
145+
'../../packages/sim-cli/dist/runtime.js',
146146
],
147147
},
148148
experimental: {

‎apps/sim/package.json‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,16 @@
88
"node": ">=22.19.0"
99
},
1010
"scripts": {
11-
"dev": "bun run build:workbench-cli && bun run dev:cache:cap && next dev --port 3000",
12-
"dev:capped": "bun run build:workbench-cli && bun run dev:cache:cap && NODE_OPTIONS='--max-old-space-size=4096' next dev --port 3000",
11+
"dev": "bun run build:cli-runtime && bun run dev:cache:cap && next dev --port 3000",
12+
"dev:capped": "bun run build:cli-runtime && bun run dev:cache:cap && NODE_OPTIONS='--max-old-space-size=4096' next dev --port 3000",
1313
"dev:cache:cap": "bun run ../../scripts/prune-turbopack-cache.ts",
1414
"dev:clean": "rm -rf .next/dev/cache",
15-
"dev:webpack": "bun run build:workbench-cli && next dev --webpack",
15+
"dev:webpack": "bun run build:cli-runtime && next dev --webpack",
1616
"load:workflow": "bun run load:workflow:baseline",
1717
"load:workflow:baseline": "BASE_URL=${BASE_URL:-http://localhost:3000} WARMUP_DURATION=${WARMUP_DURATION:-10} WARMUP_RATE=${WARMUP_RATE:-2} PEAK_RATE=${PEAK_RATE:-8} HOLD_DURATION=${HOLD_DURATION:-20} artillery run scripts/load/workflow-concurrency.yml",
1818
"load:workflow:waves": "BASE_URL=${BASE_URL:-http://localhost:3000} WAVE_ONE_DURATION=${WAVE_ONE_DURATION:-10} WAVE_ONE_RATE=${WAVE_ONE_RATE:-6} QUIET_DURATION=${QUIET_DURATION:-5} WAVE_TWO_DURATION=${WAVE_TWO_DURATION:-15} WAVE_TWO_RATE=${WAVE_TWO_RATE:-8} WAVE_THREE_DURATION=${WAVE_THREE_DURATION:-20} WAVE_THREE_RATE=${WAVE_THREE_RATE:-10} artillery run scripts/load/workflow-waves.yml",
1919
"load:workflow:isolation": "BASE_URL=${BASE_URL:-http://localhost:3000} ISOLATION_DURATION=${ISOLATION_DURATION:-30} TOTAL_RATE=${TOTAL_RATE:-9} WORKSPACE_A_WEIGHT=${WORKSPACE_A_WEIGHT:-8} WORKSPACE_B_WEIGHT=${WORKSPACE_B_WEIGHT:-1} artillery run scripts/load/workflow-isolation.yml",
20-
"build": "bun run build:workbench-cli && bun run build:sandbox-bundles && NODE_OPTIONS='--max-old-space-size=8192' next build",
21-
"build:workbench-cli": "bun build ../../packages/sim-cli/src/workbench.ts --target=node --format=esm --packages=bundle --outfile=../../packages/sim-cli/dist/workbench.js",
20+
"build": "bun run build:cli-runtime && bun run build:sandbox-bundles && NODE_OPTIONS='--max-old-space-size=8192' next build",
2221
"build:sandbox-bundles": "bun run ./lib/execution/sandbox/bundles/build.ts",
2322
"start": "next start",
2423
"prepare": "cd ../.. && bun husky",
@@ -33,7 +32,8 @@
3332
"format:check": "biome format .",
3433
"generate:pi-model-catalog": "bun run scripts/generate-pi-model-catalog.ts",
3534
"generate-docs": "bun run ../../scripts/generate-docs.ts",
36-
"org:consolidate-users": "bun run scripts/consolidate-users-into-organization.ts"
35+
"org:consolidate-users": "bun run scripts/consolidate-users-into-organization.ts",
36+
"build:cli-runtime": "bun build ../../packages/sim-cli/src/runtime.ts --target=node --format=esm --packages=bundle --outfile=../../packages/sim-cli/dist/runtime.js"
3737
},
3838
"dependencies": {
3939
"@1password/sdk": "0.3.1",

0 commit comments

Comments
 (0)