From ce775d00f66767b5d985a16d81087ac2cac2148a Mon Sep 17 00:00:00 2001 From: Parteek Singh Date: Fri, 12 Jun 2026 15:39:04 -0700 Subject: [PATCH 1/2] Fix sandbox execute file staging - Write `execute({ files })` inputs before `/execute` - Stop failed `run()` calls from waiting on output streams - Filter non-reusable storage env in Explorer demos --- docs/src/api/sandbox/route.ts | 70 +++- docs/src/api/test/sandbox-route.test.ts | 338 ++++++++++++++++++ packages/core/src/services/sandbox/execute.ts | 12 +- packages/core/src/services/sandbox/run.ts | 10 +- packages/server/test/sandbox-client.test.ts | 97 ++++- 5 files changed, 494 insertions(+), 33 deletions(-) create mode 100644 docs/src/api/test/sandbox-route.test.ts diff --git a/docs/src/api/sandbox/route.ts b/docs/src/api/sandbox/route.ts index 249299f13..bc9437fef 100644 --- a/docs/src/api/sandbox/route.ts +++ b/docs/src/api/sandbox/route.ts @@ -82,6 +82,49 @@ async function loadScriptFiles(scriptPath: string): Promise { ]; } +function firstEnvValue(...values: Array): string | undefined { + return values.find((value) => value && value.trim() !== ''); +} + +function firstReusableCredential(...values: Array): string | undefined { + return values.find((value) => value && value.trim() !== '' && !value.trim().startsWith('ags-')); +} + +function objectStorageEnv(): Record | undefined { + const bucket = firstEnvValue(process.env.S3_BUCKET, process.env.AWS_BUCKET); + const endpoint = firstEnvValue(process.env.S3_ENDPOINT, process.env.AWS_ENDPOINT); + const accessKeyId = firstReusableCredential( + process.env.S3_ACCESS_KEY_ID, + process.env.AWS_ACCESS_KEY_ID + ); + const secretAccessKey = firstReusableCredential( + process.env.S3_SECRET_ACCESS_KEY, + process.env.AWS_SECRET_ACCESS_KEY + ); + + if (!bucket || !endpoint || !accessKeyId || !secretAccessKey) return undefined; + + const storageEnv: Record = { + S3_BUCKET: bucket, + S3_ENDPOINT: endpoint, + S3_ACCESS_KEY_ID: accessKeyId, + S3_SECRET_ACCESS_KEY: secretAccessKey, + }; + const region = firstEnvValue(process.env.S3_REGION, process.env.AWS_REGION); + if (region) storageEnv.S3_REGION = region; + + return storageEnv; +} + +function copyObjectStorageEnv(envVars: Record): void { + const storageEnv = objectStorageEnv(); + if (!storageEnv) return; + + for (const [key, value] of Object.entries(storageEnv)) { + envVars[key] = value; + } +} + interface SandboxExecutionResult { readonly exitCode: number; readonly error?: string; @@ -230,23 +273,6 @@ const router = new Hono().get( envVars.AGENTUITY_CLOUD_DEPLOYMENT_ID = process.env.AGENTUITY_CLOUD_DEPLOYMENT_ID; if (process.env.DATABASE_URL) envVars.DATABASE_URL = process.env.DATABASE_URL; - const storageEnv = { - AWS_BUCKET: process.env.AWS_BUCKET ?? process.env.S3_BUCKET, - AWS_ENDPOINT: process.env.AWS_ENDPOINT ?? process.env.S3_ENDPOINT, - AWS_ACCESS_KEY_ID: process.env.AWS_ACCESS_KEY_ID ?? process.env.S3_ACCESS_KEY_ID, - AWS_SECRET_ACCESS_KEY: - process.env.AWS_SECRET_ACCESS_KEY ?? process.env.S3_SECRET_ACCESS_KEY, - AWS_REGION: process.env.AWS_REGION ?? process.env.S3_REGION, - S3_BUCKET: process.env.S3_BUCKET, - S3_ENDPOINT: process.env.S3_ENDPOINT, - S3_ACCESS_KEY_ID: process.env.S3_ACCESS_KEY_ID, - S3_SECRET_ACCESS_KEY: process.env.S3_SECRET_ACCESS_KEY, - S3_REGION: process.env.S3_REGION, - }; - for (const [key, value] of Object.entries(storageEnv)) { - if (value) envVars[key] = value; - } - const scriptPath = `dist/run/${scriptName}.js`; const command = ['bun', 'run', scriptPath, JSON.stringify(input)]; let scriptFiles: FileToWrite[]; @@ -258,8 +284,14 @@ const router = new Hono().get( await stream.writeSSE({ event: 'error', data: message }); return; } - // Sandbox env is fixed at creation, so a reused session sandbox can predate - // the linked-bucket AWS_* env (e.g. storage linked after the session began). + + if (scriptName === 'objectstore') { + // Some deployment-provided AWS_* values are scoped to the parent runtime + // and cannot be reused inside a child sandbox. Forward only reusable + // S3_* credentials for the object-storage demo. + copyObjectStorageEnv(envVars); + } + // Always run the object storage script one-shot so it sees current env. const useInteractiveSandbox = scriptName !== 'objectstore'; diff --git a/docs/src/api/test/sandbox-route.test.ts b/docs/src/api/test/sandbox-route.test.ts new file mode 100644 index 000000000..811dde013 --- /dev/null +++ b/docs/src/api/test/sandbox-route.test.ts @@ -0,0 +1,338 @@ +import { afterAll, beforeAll, expect, mock, test } from 'bun:test'; +import { mkdir, rm, writeFile } from 'node:fs/promises'; +import type { Writable } from 'node:stream'; +import type { Logger } from '@agentuity/core'; +import { KeyValueClient } from '@agentuity/keyvalue'; +import { Hono } from 'hono'; +import type { ApiEnv, ApiVariables, StateStore } from '../context'; +import { encodeSandboxOutputFrame } from '../../lib/sandbox-output-protocol'; + +Bun.env.SANDBOX_SNAPSHOT_ID = 'snapshot_test'; +Bun.env.AGENTUITY_SDK_KEY = 'sdk_test'; + +interface SandboxRunParams { + readonly options: { + readonly command: { + readonly exec: readonly string[]; + readonly files?: unknown; + }; + readonly env?: Record; + }; + readonly stdout?: Writable; +} + +interface SandboxCreateParams { + readonly options: { + readonly env?: Record; + }; +} + +interface SandboxExecuteParams { + readonly sandboxId: string; + readonly options: { + readonly command: readonly string[]; + readonly files?: unknown; + }; +} + +const sandboxRunCalls: SandboxRunParams[] = []; +const sandboxCreateCalls: SandboxCreateParams[] = []; +const sandboxExecuteCalls: SandboxExecuteParams[] = []; +const createdScriptPaths: string[] = []; +const storageEnvNames = [ + 'AWS_BUCKET', + 'AWS_ENDPOINT', + 'AWS_ACCESS_KEY_ID', + 'AWS_SECRET_ACCESS_KEY', + 'AWS_REGION', + 'S3_BUCKET', + 'S3_ENDPOINT', + 'S3_ACCESS_KEY_ID', + 'S3_SECRET_ACCESS_KEY', + 'S3_REGION', +] as const; +const silentLogger: Logger = { + debug() {}, + error() {}, + fatal(message: unknown): never { + throw new Error(String(message)); + }, + info() {}, + trace() {}, + warn() {}, + child() { + return silentLogger; + }, +}; + +const sandboxRunMock = mock(async (_client: unknown, params: SandboxRunParams) => { + sandboxRunCalls.push(params); + params.stdout?.write(encodeSandboxOutputFrame({ type: 'stdout', data: 'hello output' })); + return { sandboxId: 'sandbox_oneshot', exitCode: 0 }; +}); + +const sandboxCreateMock = mock(async (_client: unknown, params: SandboxCreateParams) => { + sandboxCreateCalls.push(params); + return { sandboxId: 'sandbox_interactive' }; +}); + +const sandboxExecuteMock = mock(async (_client: unknown, params: SandboxExecuteParams) => { + sandboxExecuteCalls.push(params); + return { executionId: 'execution_1' }; +}); + +const executionGetMock = mock(async () => ({ + status: 'completed', + exitCode: 0, +})); + +class APIClientMock {} +class SandboxNotFoundErrorMock extends Error {} +class SandboxTerminatedErrorMock extends Error {} + +mock.module('@agentuity/server', () => ({ + APIClient: APIClientMock, + SandboxNotFoundError: SandboxNotFoundErrorMock, + SandboxTerminatedError: SandboxTerminatedErrorMock, + executionGet: executionGetMock, + getServiceUrls: () => ({ sandbox: 'https://sandbox.test' }), + sandboxCreate: sandboxCreateMock, + sandboxExecute: sandboxExecuteMock, + sandboxRun: sandboxRunMock, +})); + +const { default: sandboxRouter } = await import('../sandbox/route'); + +async function ensureRunScript(scriptName: string): Promise { + const scriptPath = `dist/run/${scriptName}.js`; + if (await Bun.file(scriptPath).exists()) return; + + await mkdir('dist/run', { recursive: true }); + await writeFile(scriptPath, 'console.log("test script");\n'); + createdScriptPaths.push(scriptPath); +} + +function storageEnvKeys(env: Record | undefined): string[] { + return Object.keys(env ?? {}) + .filter((key) => key.startsWith('AWS_') || key.startsWith('S3_')) + .sort(); +} + +function snapshotStorageEnv(): Record<(typeof storageEnvNames)[number], string | undefined> { + return Object.fromEntries(storageEnvNames.map((key) => [key, process.env[key]])) as Record< + (typeof storageEnvNames)[number], + string | undefined + >; +} + +function setStorageEnv( + values: Partial> +): void { + for (const key of storageEnvNames) { + const value = values[key]; + if (value === undefined) { + delete process.env[key]; + continue; + } + process.env[key] = value; + } +} + +beforeAll(async () => { + await Promise.all([ensureRunScript('hello'), ensureRunScript('objectstore')]); +}); + +afterAll(async () => { + for (const scriptPath of createdScriptPaths) { + await rm(scriptPath, { force: true }); + } +}); + +function createLogger(): ApiVariables['logger'] { + return silentLogger; +} + +function createKv(): ApiVariables['kv'] { + const kv = new KeyValueClient({ + apiKey: 'sdk_test', + logger: createLogger(), + url: 'https://keyvalue.test', + }); + Object.defineProperty(kv, 'get', { + value: async () => ({ exists: false }), + }); + Object.defineProperty(kv, 'set', { + value: async () => {}, + }); + return kv; +} + +function createStateStore(): StateStore { + return { + async delete() {}, + async entries() { + return []; + }, + async get(): Promise { + return undefined; + }, + async has() { + return false; + }, + async push() {}, + async set() {}, + }; +} + +function createThread(): ApiVariables['thread'] { + return { + id: 'thread_test', + state: createStateStore(), + }; +} + +function createApp(options: { readonly interactive: boolean }): Hono { + const app = new Hono(); + app.use('*', async (c, next) => { + c.set('logger', createLogger()); + if (options.interactive) { + c.set('kv', createKv()); + c.set('thread', createThread()); + } + await next(); + }); + app.route('/api/sandbox', sandboxRouter); + return app; +} + +test('sandbox route completes one-shot execution without storage credentials for non-storage scripts', async () => { + sandboxRunCalls.length = 0; + + const response = await createApp({ interactive: false }).fetch( + new Request('http://docs.test/api/sandbox/run?script=hello') + ); + const body = await response.text(); + + expect(response.status).toBe(200); + expect(body).toContain('event: stdout'); + expect(body).toContain('data: hello output'); + expect(body).toContain('event: done'); + expect(sandboxRunCalls).toHaveLength(1); + expect(sandboxRunCalls[0]?.options.command.files).toHaveLength(1); + expect(storageEnvKeys(sandboxRunCalls[0]?.options.env)).toEqual([]); +}); + +test('sandbox route completes interactive execution without storage credentials for non-storage scripts', async () => { + sandboxCreateCalls.length = 0; + sandboxExecuteCalls.length = 0; + + const response = await createApp({ interactive: true }).fetch( + new Request('http://docs.test/api/sandbox/run?script=hello') + ); + const body = await response.text(); + + expect(response.status).toBe(200); + expect(body).toContain('event: status'); + expect(body).toContain('event: done'); + expect(sandboxCreateCalls).toHaveLength(1); + expect(sandboxExecuteCalls).toHaveLength(1); + expect(storageEnvKeys(sandboxCreateCalls[0]?.options.env)).toEqual([]); + expect(sandboxExecuteCalls[0]?.options.files).toHaveLength(1); +}); + +test('sandbox route passes object storage credentials as S3 env only', async () => { + const previousEnv = snapshotStorageEnv(); + sandboxRunCalls.length = 0; + setStorageEnv({ + AWS_BUCKET: 'bucket_test', + AWS_ENDPOINT: 'https://storage.test', + AWS_ACCESS_KEY_ID: 'access_test', + AWS_SECRET_ACCESS_KEY: 'secret_test', + AWS_REGION: 'auto', + }); + + try { + const response = await createApp({ interactive: false }).fetch( + new Request('http://docs.test/api/sandbox/run?script=objectstore') + ); + const body = await response.text(); + + expect(response.status).toBe(200); + expect(body).toContain('event: done'); + expect(sandboxRunCalls).toHaveLength(1); + expect(sandboxRunCalls[0]?.options.env?.S3_BUCKET).toBe('bucket_test'); + expect(sandboxRunCalls[0]?.options.env?.S3_ENDPOINT).toBe('https://storage.test'); + expect(sandboxRunCalls[0]?.options.env?.S3_ACCESS_KEY_ID).toBe('access_test'); + expect(sandboxRunCalls[0]?.options.env?.S3_SECRET_ACCESS_KEY).toBe('secret_test'); + expect(sandboxRunCalls[0]?.options.env?.S3_REGION).toBe('auto'); + expect( + storageEnvKeys(sandboxRunCalls[0]?.options.env).filter((key) => key.startsWith('AWS_')) + ).toEqual([]); + } finally { + setStorageEnv(previousEnv); + } +}); + +test('sandbox route prefers reusable S3 credentials over hashed AWS aliases', async () => { + const previousEnv = snapshotStorageEnv(); + sandboxRunCalls.length = 0; + setStorageEnv({ + AWS_BUCKET: 'aws_bucket', + AWS_ENDPOINT: 'https://aws-storage.test', + AWS_ACCESS_KEY_ID: 'ags-access-token', + AWS_SECRET_ACCESS_KEY: 'ags-secret-token', + AWS_REGION: 'usc', + S3_BUCKET: 's3_bucket', + S3_ENDPOINT: 'https://s3-storage.test', + S3_ACCESS_KEY_ID: 's3_access', + S3_SECRET_ACCESS_KEY: 's3_secret', + S3_REGION: 'auto', + }); + + try { + const response = await createApp({ interactive: false }).fetch( + new Request('http://docs.test/api/sandbox/run?script=objectstore') + ); + const body = await response.text(); + + expect(response.status).toBe(200); + expect(body).toContain('event: done'); + expect(sandboxRunCalls).toHaveLength(1); + expect(sandboxRunCalls[0]?.options.env?.S3_BUCKET).toBe('s3_bucket'); + expect(sandboxRunCalls[0]?.options.env?.S3_ENDPOINT).toBe('https://s3-storage.test'); + expect(sandboxRunCalls[0]?.options.env?.S3_ACCESS_KEY_ID).toBe('s3_access'); + expect(sandboxRunCalls[0]?.options.env?.S3_SECRET_ACCESS_KEY).toBe('s3_secret'); + expect(sandboxRunCalls[0]?.options.env?.S3_REGION).toBe('auto'); + expect( + storageEnvKeys(sandboxRunCalls[0]?.options.env).filter((key) => key.startsWith('AWS_')) + ).toEqual([]); + } finally { + setStorageEnv(previousEnv); + } +}); + +test('sandbox route does not forward hashed object storage credentials', async () => { + const previousEnv = snapshotStorageEnv(); + sandboxRunCalls.length = 0; + setStorageEnv({ + AWS_BUCKET: 'bucket_test', + AWS_ENDPOINT: 'https://storage.test', + AWS_ACCESS_KEY_ID: 'ags-access-token', + AWS_SECRET_ACCESS_KEY: 'ags-secret-token', + AWS_REGION: 'auto', + }); + + try { + const response = await createApp({ interactive: false }).fetch( + new Request('http://docs.test/api/sandbox/run?script=objectstore') + ); + const body = await response.text(); + + expect(response.status).toBe(200); + expect(body).toContain('event: done'); + expect(sandboxRunCalls).toHaveLength(1); + expect(storageEnvKeys(sandboxRunCalls[0]?.options.env)).toEqual([]); + } finally { + setStorageEnv(previousEnv); + } +}); diff --git a/packages/core/src/services/sandbox/execute.ts b/packages/core/src/services/sandbox/execute.ts index 49814544b..74035e737 100644 --- a/packages/core/src/services/sandbox/execute.ts +++ b/packages/core/src/services/sandbox/execute.ts @@ -2,7 +2,7 @@ import type { ExecuteOptions, Execution, ExecutionStatus } from './types.ts'; import { z } from 'zod'; import type { APIClient } from '../api.ts'; import { SandboxBusyError, SandboxNotFoundError, throwSandboxError } from './util.ts'; -import { base64Encode } from './base64.ts'; +import { sandboxWriteFiles } from './files.ts'; export const ExecuteRequestSchema = z .object({ @@ -86,10 +86,12 @@ export async function sandboxExecute( }; if (options.files && options.files.length > 0) { - body.files = options.files.map((f) => ({ - path: f.path, - content: base64Encode(f.content), - })); + await sandboxWriteFiles(client, { + sandboxId, + files: options.files, + orgId, + signal: signal ?? options.signal, + }); } if (options.timeout) { body.timeout = options.timeout; diff --git a/packages/core/src/services/sandbox/run.ts b/packages/core/src/services/sandbox/run.ts index ddd1a76de..712fce0d6 100644 --- a/packages/core/src/services/sandbox/run.ts +++ b/packages/core/src/services/sandbox/run.ts @@ -253,10 +253,14 @@ export async function sandboxRun( ); try { - // Resolve execution/status first so a hung Pulse reader cannot block the - // whole run until the client deadline. Output stream fetches still run in parallel. - const [execution] = await Promise.all([completionPromise, streamsPromise]); + // Resolve execution/status first so a failed create is not hidden by + // Pulse streams that stay open until the client deadline. + const execution = await completionPromise; finalExecution = execution; + if (execution.status !== 'completed') { + abortController.abort(); + } + await streamsPromise; abortController.abort(); } catch (error) { throw mapRunAbortError( diff --git a/packages/server/test/sandbox-client.test.ts b/packages/server/test/sandbox-client.test.ts index a2a4fc725..a419bc668 100644 --- a/packages/server/test/sandbox-client.test.ts +++ b/packages/server/test/sandbox-client.test.ts @@ -258,12 +258,24 @@ describe('SandboxClient', () => { expect(result.exitCode).toBe(0); }); - test('execute with files should send files as map in request body', async () => { - let requestBody: Record | null = null; + test('execute with files should stage files before executing command', async () => { + let executeRequestBody: Record | null = null; + let writeFilesRequestBody: Record | null = null; mockFetch(async (url, opts) => { + if (opts?.method === 'POST' && url.includes('/fs/sandbox-123')) { + writeFilesRequestBody = JSON.parse(opts.body as string); + return new Response( + JSON.stringify({ + success: true, + data: { filesWritten: 2 }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + if (opts?.method === 'POST' && url.includes('/execute')) { - requestBody = JSON.parse(opts.body as string); + executeRequestBody = JSON.parse(opts.body as string); return new Response( JSON.stringify({ success: true, @@ -315,12 +327,14 @@ describe('SandboxClient', () => { ], }); - expect(requestBody).not.toBeNull(); - expect(requestBody!.command).toEqual(['bun', 'run', 'script.ts']); - expect(requestBody!.files).toEqual([ + expect(writeFilesRequestBody).not.toBeNull(); + expect(writeFilesRequestBody!.files).toEqual([ { path: 'script.ts', content: Buffer.from('console.log("hello")').toString('base64') }, { path: 'data.json', content: Buffer.from('{"key": "value"}').toString('base64') }, ]); + expect(executeRequestBody).not.toBeNull(); + expect(executeRequestBody!.command).toEqual(['bun', 'run', 'script.ts']); + expect(executeRequestBody!.files).toBeUndefined(); }); test('execute with empty files array should not include files in request', async () => { @@ -997,6 +1011,77 @@ describe('SandboxClient', () => { expect(result.exitCode).toBe(1); }); + test('should return failed execution result without waiting for hung output streams', async () => { + let executionPolls = 0; + + mockFetch(async (url, opts) => { + if (opts?.method === 'POST' && url.includes('/sandbox')) { + return new Response( + JSON.stringify({ + success: true, + data: { + sandboxId: 'sandbox-failed-stream', + executionId: 'exec-failed-stream', + status: 'running', + stdoutStreamUrl: 'https://stream.example.com/combined/failed-stream', + stderrStreamUrl: 'https://stream.example.com/combined/failed-stream', + }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + if (url.includes('stream.example.com/combined/failed-stream')) { + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode('booting...\n')); + opts?.signal?.addEventListener('abort', () => controller.close(), { + once: true, + }); + }, + }), + { status: 200 } + ); + } + + if (url.includes('/execution/exec-failed-stream')) { + executionPolls++; + return new Response( + JSON.stringify({ + success: true, + data: { + executionId: 'exec-failed-stream', + sandboxId: 'sandbox-failed-stream', + status: 'failed', + error: 'error creating sandbox', + }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + return new Response(null, { status: 404 }); + }); + + const client = new SandboxClient({ logger: createMockLogger() }); + const abortController = new AbortController(); + const timeout = setTimeout(() => abortController.abort(), 100); + + try { + const result = await client.run( + { command: { exec: ['false'] } }, + { signal: abortController.signal } + ); + + expect(result.sandboxId).toBe('sandbox-failed-stream'); + expect(result.exitCode).toBe(1); + expect(executionPolls).toBe(1); + } finally { + clearTimeout(timeout); + } + }); + test('should return captured stdout in result', async () => { const stdoutChunks = [ new Uint8Array([72, 101, 108, 108, 111]), // "Hello" From 5daa8b301f1e67ca59fcebbb701eb1a15f29ae70 Mon Sep 17 00:00:00 2001 From: Parteek Singh Date: Fri, 12 Jun 2026 15:51:29 -0700 Subject: [PATCH 2/2] Address sandbox review follow-ups - Roll back staged execute files when `/execute` is rejected - Abort the losing `run()` completion waiter after a race win - Cover both paths in sandbox client tests --- packages/core/src/services/sandbox/execute.ts | 55 ++++++- packages/core/src/services/sandbox/run.ts | 1 + packages/server/test/sandbox-client.test.ts | 136 ++++++++++++++++++ 3 files changed, 190 insertions(+), 2 deletions(-) diff --git a/packages/core/src/services/sandbox/execute.ts b/packages/core/src/services/sandbox/execute.ts index 74035e737..4850cfc01 100644 --- a/packages/core/src/services/sandbox/execute.ts +++ b/packages/core/src/services/sandbox/execute.ts @@ -1,8 +1,10 @@ -import type { ExecuteOptions, Execution, ExecutionStatus } from './types.ts'; +import type { ExecuteOptions, Execution, ExecutionStatus, FileToWrite } from './types.ts'; import { z } from 'zod'; import type { APIClient } from '../api.ts'; import { SandboxBusyError, SandboxNotFoundError, throwSandboxError } from './util.ts'; -import { sandboxWriteFiles } from './files.ts'; +import { sandboxRmFile, sandboxWriteFiles } from './files.ts'; + +const EXECUTE_FILE_ROLLBACK_TIMEOUT_MS = 5_000; export const ExecuteRequestSchema = z .object({ @@ -68,6 +70,47 @@ export const SandboxExecuteParamsSchema = z.object({ export type SandboxExecuteParams = z.infer; +function createRollbackSignal(): { readonly signal: AbortSignal; readonly cleanup: () => void } { + if (typeof AbortSignal.timeout === 'function') { + return { signal: AbortSignal.timeout(EXECUTE_FILE_ROLLBACK_TIMEOUT_MS), cleanup() {} }; + } + + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), EXECUTE_FILE_ROLLBACK_TIMEOUT_MS); + return { + signal: controller.signal, + cleanup() { + clearTimeout(timer); + }, + }; +} + +async function cleanupStagedExecuteFiles( + client: APIClient, + params: { + readonly sandboxId: string; + readonly files: readonly FileToWrite[]; + readonly orgId: string | undefined; + } +): Promise { + const { sandboxId, files, orgId } = params; + const rollback = createRollbackSignal(); + try { + await Promise.allSettled( + files.map((file) => + sandboxRmFile(client, { + sandboxId, + path: file.path, + orgId, + signal: rollback.signal, + }) + ) + ); + } finally { + rollback.cleanup(); + } +} + /** * Executes a command in an existing sandbox. * @@ -84,6 +127,7 @@ export async function sandboxExecute( const body: z.infer = { command: options.command, }; + let stagedFiles: readonly FileToWrite[] | undefined; if (options.files && options.files.length > 0) { await sandboxWriteFiles(client, { @@ -92,6 +136,7 @@ export async function sandboxExecute( orgId, signal: signal ?? options.signal, }); + stagedFiles = options.files; } if (options.timeout) { body.timeout = options.timeout; @@ -117,6 +162,9 @@ export async function sandboxExecute( signal ?? options.signal ); } catch (error: unknown) { + if (stagedFiles) { + await cleanupStagedExecuteFiles(client, { sandboxId, files: stagedFiles, orgId }); + } if ( error && typeof error === 'object' && @@ -165,5 +213,8 @@ export async function sandboxExecute( }; } + if (stagedFiles) { + await cleanupStagedExecuteFiles(client, { sandboxId, files: stagedFiles, orgId }); + } throwSandboxError(resp, { sandboxId }); } diff --git a/packages/core/src/services/sandbox/run.ts b/packages/core/src/services/sandbox/run.ts index 712fce0d6..6b6bf0ee3 100644 --- a/packages/core/src/services/sandbox/run.ts +++ b/packages/core/src/services/sandbox/run.ts @@ -515,6 +515,7 @@ async function waitForRunCompletion( const result = await Promise.race([executionPromise, statusPromise]); return result; } finally { + completionAbortController.abort(); if (onAbort && signal) { signal.removeEventListener('abort', onAbort); } diff --git a/packages/server/test/sandbox-client.test.ts b/packages/server/test/sandbox-client.test.ts index a419bc668..c52529a0b 100644 --- a/packages/server/test/sandbox-client.test.ts +++ b/packages/server/test/sandbox-client.test.ts @@ -337,6 +337,68 @@ describe('SandboxClient', () => { expect(executeRequestBody!.files).toBeUndefined(); }); + test('execute with files should remove staged files when execute is rejected', async () => { + const removedPaths: string[] = []; + + mockFetch(async (url, opts) => { + if (opts?.method === 'POST' && url.includes('/fs/rm/sandbox-123')) { + const body = JSON.parse(opts.body as string) as { path: string }; + removedPaths.push(body.path); + return new Response(JSON.stringify({ success: true, found: true }), { + status: 200, + headers: { 'content-type': 'application/json' }, + }); + } + + if (opts?.method === 'POST' && url.includes('/fs/sandbox-123')) { + return new Response( + JSON.stringify({ + success: true, + data: { filesWritten: 2 }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + if (opts?.method === 'POST' && url.includes('/execute')) { + return new Response( + JSON.stringify({ + success: false, + message: 'sandbox is busy', + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + if (opts?.method === 'POST' && url.includes('/sandbox')) { + return new Response( + JSON.stringify({ + success: true, + data: { sandboxId: 'sandbox-123', status: 'idle' }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + return new Response(null, { status: 404 }); + }); + + const client = new SandboxClient({ logger: createMockLogger() }); + const sandbox = await client.create(); + + await expect( + sandbox.execute({ + command: ['bun', 'run', 'script.ts'], + files: [ + { path: 'script.ts', content: Buffer.from('console.log("hello")') }, + { path: 'data.json', content: Buffer.from('{"key": "value"}') }, + ], + }) + ).rejects.toThrow('sandbox is busy'); + + expect(removedPaths.sort()).toEqual(['data.json', 'script.ts']); + }); + test('execute with empty files array should not include files in request', async () => { let requestBody: Record | null = null; @@ -1082,6 +1144,80 @@ describe('SandboxClient', () => { } }); + test('should abort status completion waiter after execution completion wins', async () => { + let resolveStatusStarted: (() => void) | undefined; + const statusStarted = new Promise((resolve) => { + resolveStatusStarted = resolve; + }); + let statusWaitAborted = false; + + mockFetch(async (url, opts) => { + if (opts?.method === 'POST' && url.includes('/sandbox')) { + return new Response( + JSON.stringify({ + success: true, + data: { + sandboxId: 'sandbox-race-test', + executionId: 'exec-race-test', + status: 'running', + }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + if (url.includes('/sandbox/status/sandbox-race-test')) { + resolveStatusStarted?.(); + return new Promise((resolve) => { + opts?.signal?.addEventListener( + 'abort', + () => { + statusWaitAborted = true; + resolve( + new Response( + JSON.stringify({ + success: true, + data: { + sandboxId: 'sandbox-race-test', + status: 'running', + }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ) + ); + }, + { once: true } + ); + }); + } + + if (url.includes('/execution/exec-race-test')) { + await statusStarted; + return new Response( + JSON.stringify({ + success: true, + data: { + executionId: 'exec-race-test', + sandboxId: 'sandbox-race-test', + status: 'completed', + exitCode: 0, + }, + }), + { status: 200, headers: { 'content-type': 'application/json' } } + ); + } + + return new Response(null, { status: 404 }); + }); + + const client = new SandboxClient({ logger: createMockLogger() }); + const result = await client.run({ command: { exec: ['true'] } }); + + expect(result.sandboxId).toBe('sandbox-race-test'); + expect(result.exitCode).toBe(0); + expect(statusWaitAborted).toBe(true); + }); + test('should return captured stdout in result', async () => { const stdoutChunks = [ new Uint8Array([72, 101, 108, 108, 111]), // "Hello"