diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d73538a..29f9105 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -33,15 +33,15 @@ jobs: - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2 with: - bun-version: 1.3.13 + bun-version: 1.3.14 - name: Cache Bun packages uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5 with: path: ~/.bun/install/cache - key: ${{ runner.os }}-codeapi-api-bun-1.3.13-${{ hashFiles('api/bun.lock') }} + key: ${{ runner.os }}-codeapi-api-bun-1.3.14-${{ hashFiles('api/bun.lock') }} restore-keys: | - ${{ runner.os }}-codeapi-api-bun-1.3.13- + ${{ runner.os }}-codeapi-api-bun-1.3.14- - name: Install dependencies run: bun ci @@ -82,18 +82,108 @@ jobs: - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2 with: - bun-version: 1.3.13 + bun-version: 1.3.14 - name: Cache Bun packages uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5 with: path: ~/.bun/install/cache - key: ${{ runner.os }}-codeapi-service-bun-1.3.13-${{ hashFiles('service/bun.lock') }} + key: ${{ runner.os }}-codeapi-service-bun-1.3.14-${{ hashFiles('service/bun.lock') }} restore-keys: | - ${{ runner.os }}-codeapi-service-bun-1.3.13- + ${{ runner.os }}-codeapi-service-bun-1.3.14- - name: Install dependencies run: bun ci + - name: Build service + run: bun run build + - name: Bun tests run: bun run test + + lambda-microvm-provisioning: + name: Lambda MicroVM Provisioning + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + + - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2 + with: + bun-version: 1.3.14 + + - name: Install service dependencies + working-directory: service + run: bun ci + + - name: Type-check MicroVM image helper + working-directory: service + run: >- + bunx tsc --noEmit + --moduleResolution node16 + --module node16 + --target es2022 + --strict + --esModuleInterop + --skipLibCheck + --types node,bun-types + scripts/create-microvm-image.ts + + - name: Test MicroVM image helper state machine + working-directory: service + run: bun test scripts/create-microvm-image.test.ts + + - name: Validate artifact build script + run: | + bash -n scripts/build-lambda-microvm-artifact.sh + shellcheck scripts/build-lambda-microvm-artifact.sh + + - name: Validate runner Dockerfile + run: >- + docker buildx build --check + --platform linux/arm64 + --target lambda-microvm-runner + -f api/Dockerfile + . + + - uses: hashicorp/setup-terraform@dfe3c3f87815947d99a8997f908cb6525fc44e9e # v4.0.1 + with: + terraform_version: 1.15.2 + terraform_wrapper: false + + - name: Validate Terraform + working-directory: docs/lambda-microvm/terraform + run: | + terraform fmt -check -recursive + terraform init -backend=false -input=false -lockfile=readonly + terraform validate + + lambda-microvm-runner-build: + name: Lambda MicroVM Runner Image (arm64) + needs: lambda-microvm-provisioning + runs-on: ubuntu-24.04-arm + timeout-minutes: 60 + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + + - name: Create native BuildKit builder + run: | + test "$(uname -m)" = "aarch64" + docker buildx create \ + --name lambda-microvm-ci \ + --driver docker-container \ + --use + docker buildx inspect --bootstrap + + - name: Build Lambda MicroVM runner image + # The target pins Bun 1.3.14 and installs api/bun.lock frozen. Building + # natively catches architecture-specific package and Docker layer + # failures without retaining a second multi-gigabyte image copy. + run: | + docker buildx build \ + --builder lambda-microvm-ci \ + --platform linux/arm64 \ + --target lambda-microvm-runner \ + --output type=cacheonly \ + --progress plain \ + -f api/Dockerfile \ + . diff --git a/.gitignore b/.gitignore index 36b678d..db2b8a2 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ node_modules /file-server-logs/**/* /service-logs/**/* /tool-call-server-logs/**/* +/.build-lambda-microvm/ # Helm artifacts helm/*/charts/*.tgz diff --git a/api/Dockerfile b/api/Dockerfile index ff4e848..9cc7525 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -121,8 +121,8 @@ COPY --from=node:22-slim /usr/local/bin/node /usr/local/bin/node WORKDIR /sandbox_api -COPY api/package.json ./ -RUN bun install +COPY api/package.json api/bun.lock ./ +RUN bun install --frozen-lockfile COPY api/tsconfig.json ./ COPY api/src ./src @@ -148,6 +148,27 @@ RUN rm -f /usr/bin/nsenter /usr/bin/unshare /usr/bin/chroot /usr/sbin/chroot \ COPY api/src/entrypoint.sh ./entrypoint.sh RUN chmod +x ./entrypoint.sh +# ============================================================================ +# Stage 2b: AWS Lambda MicroVM container base image +# +# Lambda MicroVM IS the VM boundary: no libkrun launcher, no guest kernel. +# This target is pushed to a same-account ECR repo and referenced by the +# FROM line of the code-artifact Dockerfile that Lambda builds on top of +# the AL2023 MicroVM base image (see scripts/build-lambda-microvm-artifact.sh). +# Inbound traffic defaults to port 8080; lifecycle/build hooks are served +# by the same app at /aws/lambda-microvms/runtime/v1/*. +# ============================================================================ +FROM sandbox-build AS lambda-microvm-runner + +COPY --from=package-builder /pkgs /pkgs + +ENV PORT=8080 \ + SANDBOX_PACKAGES_DIRECTORY=/pkgs \ + SANDBOX_SESSION_WORKSPACE_ENABLED=true + +EXPOSE 8080/tcp +ENTRYPOINT ["/sandbox_api/entrypoint.sh"] + # ============================================================================ # Stage 3: Build the Rust launcher binary (Fedora for libkrun ABI) # ============================================================================ diff --git a/api/README.md b/api/README.md index 9745b73..07c437c 100644 --- a/api/README.md +++ b/api/README.md @@ -40,7 +40,7 @@ All prefixed with `SANDBOX_` unless noted: | `SANDBOX_PACKAGES_DIRECTORY` | `/pkgs` | Directory containing language packages | | `SANDBOX_DISABLE_NETWORKING` | `true` | Isolate sandbox from the network | | `SANDBOX_ALLOWED_LOCAL_NETWORK_PORT` | `0` | Allow sandbox to reach this host port (for tool calling) | -| `SANDBOX_OUTPUT_MAX_SIZE` | `1024` | Max stdout/stderr bytes before truncation | +| `SANDBOX_OUTPUT_MAX_SIZE` | `1024` | Per-stream output cap. stderr truncates at this size (the job keeps running); stdout overflow kills the job (`status: OL`), since stdout is the result. Every shipped compose/helm config sets `65536` — the `1024` fallback only applies when running the runner bare. | | `SANDBOX_MAX_PROCESS_COUNT` | `64` | Max PIDs inside the sandbox | | `SANDBOX_MAX_OPEN_FILES` | `2048` | rlimit nofile | | `SANDBOX_MAX_FILE_SIZE` | `10000000` | rlimit fsize (bytes) | diff --git a/api/src/api/lifecycle.test.ts b/api/src/api/lifecycle.test.ts new file mode 100644 index 0000000..dcf4bbf --- /dev/null +++ b/api/src/api/lifecycle.test.ts @@ -0,0 +1,59 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { + LIFECYCLE_HOOK_BASE_PATH, + applyRunHook, + getMicrovmRunContext, + resetMicrovmRunContextForTests, +} from './lifecycle'; + +afterEach(resetMicrovmRunContextForTests); + +describe('MicroVM /run hook context', () => { + test('captures microvmId and runHookPayload from the platform body', () => { + const context = applyRunHook({ + microvmId: 'mvm-0123', + runHookPayload: '{"runtime_session_id":"rt_x"}', + }); + expect(context.microvmId).toBe('mvm-0123'); + expect(context.runHookPayload).toBe('{"runtime_session_id":"rt_x"}'); + expect(getMicrovmRunContext()).toBe(context); + }); + + test('first run wins: a different microvmId does not overwrite the context', () => { + applyRunHook({ microvmId: 'mvm-first', runHookPayload: 'a' }); + applyRunHook({ microvmId: 'mvm-second', runHookPayload: 'b' }); + expect(getMicrovmRunContext()?.microvmId).toBe('mvm-first'); + expect(getMicrovmRunContext()?.runHookPayload).toBe('a'); + }); + + test('retries with the same microvmId are idempotent', () => { + const first = applyRunHook({ microvmId: 'mvm-0123' }); + const second = applyRunHook({ microvmId: 'mvm-0123' }); + expect(second).toBe(first); + }); + + test('tolerates malformed and empty bodies', () => { + expect(applyRunHook(undefined).microvmId).toBeUndefined(); + resetMicrovmRunContextForTests(); + expect(applyRunHook('not-an-object').microvmId).toBeUndefined(); + resetMicrovmRunContextForTests(); + expect(applyRunHook({ microvmId: 42, runHookPayload: {} }).microvmId).toBeUndefined(); + }); + + test('a malformed first delivery does not block a valid retry', () => { + const malformed = applyRunHook({ runHookPayload: '{"runtime_session_id":"rt_bad"}' }); + expect(malformed.microvmId).toBeUndefined(); + expect(getMicrovmRunContext()).toBeUndefined(); + + const retried = applyRunHook({ + microvmId: 'mvm-retried', + runHookPayload: '{"runtime_session_id":"rt_good"}', + }); + expect(retried.microvmId).toBe('mvm-retried'); + expect(getMicrovmRunContext()).toBe(retried); + }); + + test('hook base path matches the AWS well-known prefix', () => { + expect(LIFECYCLE_HOOK_BASE_PATH).toBe('/aws/lambda-microvms/runtime/v1'); + }); +}); diff --git a/api/src/api/lifecycle.ts b/api/src/api/lifecycle.ts new file mode 100644 index 0000000..28977ed --- /dev/null +++ b/api/src/api/lifecycle.ts @@ -0,0 +1,104 @@ +import express, { Router, type Request, type Response } from 'express'; +import { logger } from '../logger'; +import { bindSessionWorkspace, parseSessionBinding, unbindSessionWorkspace } from '../session-workspace'; + +/** + * AWS Lambda MicroVM hook endpoints. The platform POSTs to + * `/aws/lambda-microvms/runtime/v1/` on the configured hook port: + * + * /ready, /validate image build hooks (200 = proceed, 503 = retry) + * /run after start; body {microvmId, runHookPayload}; + * external traffic only flows after this returns 200 + * /resume, /suspend, /terminate lifecycle transitions + * + * Phase 1-2: structured no-ops that log and 200. /suspend and /terminate + * are the Phase 3 checkpoint-flush attachment points; /run captures the + * per-VM payload for later runtime-session wiring. + */ + +export const LIFECYCLE_HOOK_BASE_PATH = '/aws/lambda-microvms/runtime/v1'; + +export interface MicrovmRunContext { + microvmId?: string; + runHookPayload?: string; + receivedAt: number; +} + +let runContext: MicrovmRunContext | undefined; + +export function getMicrovmRunContext(): MicrovmRunContext | undefined { + return runContext; +} + +export function resetMicrovmRunContextForTests(): void { + runContext = undefined; +} + +/** First /run wins; retries with the same microvmId are idempotent, a + * different id is logged and ignored. Always 200 — a non-200 would fail + * the platform's lifecycle operation. */ +export function applyRunHook(body: unknown): MicrovmRunContext { + const parsed = (typeof body === 'object' && body !== null ? body : {}) as { + microvmId?: unknown; + runHookPayload?: unknown; + }; + const microvmId = typeof parsed.microvmId === 'string' ? parsed.microvmId : undefined; + const runHookPayload = typeof parsed.runHookPayload === 'string' ? parsed.runHookPayload : undefined; + + /* A malformed first delivery must not permanently win the idempotency slot. + * Lambda can retry lifecycle hooks, so only a payload carrying the platform + * identity becomes the immutable run context. Keep returning an ephemeral + * context for malformed calls because the hook must still acknowledge them + * with 200 rather than failing the lifecycle operation. */ + if (microvmId == null) { + logger.warn('Ignoring /run hook without a valid microvmId'); + return runContext ?? { microvmId, runHookPayload, receivedAt: Date.now() }; + } + + if (runContext == null) { + runContext = { microvmId, runHookPayload, receivedAt: Date.now() }; + const binding = parseSessionBinding(runHookPayload); + if (binding) { + bindSessionWorkspace(binding); + logger.info({ runtimeSessionId: binding.runtimeSessionId }, 'Bound persistent session workspace'); + } + } else if (runContext.microvmId != null && microvmId != null && runContext.microvmId !== microvmId) { + logger.warn( + { existing: runContext.microvmId, incoming: microvmId }, + 'Ignoring /run hook for a different microvmId', + ); + } + return runContext; +} + +const lifecycleRouter = Router(); + +function ackHook(hook: string) { + return (_req: Request, res: Response): Response => { + logger.info({ hook }, 'MicroVM lifecycle hook invoked'); + return res.status(200).json({ hook, status: 'ok' }); + }; +} + +lifecycleRouter.post('/ready', ackHook('ready')); +lifecycleRouter.post('/validate', ackHook('validate')); + +lifecycleRouter.post('/run', express.json({ limit: '32kb' }), (req: Request, res: Response) => { + const context = applyRunHook(req.body); + logger.info( + { hook: 'run', microvmId: context.microvmId, hasPayload: context.runHookPayload != null }, + 'MicroVM lifecycle hook invoked', + ); + return res.status(200).json({ hook: 'run', status: 'ok' }); +}); + +lifecycleRouter.post('/resume', ackHook('resume')); +lifecycleRouter.post('/suspend', ackHook('suspend')); + +lifecycleRouter.post('/terminate', (_req: Request, res: Response) => { + logger.info({ hook: 'terminate' }, 'MicroVM lifecycle hook invoked'); + void unbindSessionWorkspace().catch((err) => logger.error({ err }, 'Failed to unbind session workspace on terminate')); + return res.status(200).json({ hook: 'terminate', status: 'ok' }); +}); + +export default lifecycleRouter; diff --git a/api/src/api/session-inputs.routes.test.ts b/api/src/api/session-inputs.routes.test.ts new file mode 100644 index 0000000..48fe7d5 --- /dev/null +++ b/api/src/api/session-inputs.routes.test.ts @@ -0,0 +1,132 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from 'bun:test'; +import { spawnSync } from 'child_process'; +import express from 'express'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import type { Server } from 'http'; +import v2Router from './v2'; +import { config } from '../config'; +import { + SESSION_INPUT_CACHE_DIR, + SESSION_INPUT_CACHE_MAX_OBJECTS, + inputCacheKey, +} from '../session-inputs'; + +/** + * Route-level coverage for input delivery. The unit suites exercise the cache + * itself; this exercises the WIRING — body parsing, the JSON gate's tar + * exemption, and mount paths — which is where a live-only failure hid: the + * probe route had no parser (there is no global one), so every ref list came + * back "refs must be an array". + */ + +let server: Server; +let baseUrl: string; +const savedSessionWorkspaceEnabled = config.session_workspace_enabled; + +beforeAll(async () => { + config.session_workspace_enabled = true; + const app = express(); + /* Mirror index.ts: no global JSON parser, router mounted under /api/v2. */ + app.use(express.urlencoded({ extended: true })); + app.use('/api/v2', v2Router); + await new Promise((resolve) => { + server = app.listen(0, '127.0.0.1', () => resolve()); + }); + const address = server.address(); + baseUrl = `http://127.0.0.1:${typeof address === 'object' && address ? address.port : 0}`; +}); + +afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); + config.session_workspace_enabled = savedSessionWorkspaceEnabled; +}); + +afterEach(async () => { + config.session_workspace_enabled = true; + await fsp.rm(SESSION_INPUT_CACHE_DIR, { recursive: true, force: true }).catch(() => {}); +}); + +async function makeBatch(entries: Array<{ sid: string; id: string; body: string }>): Promise { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'route-batch-')); + for (const entry of entries) { + const key = inputCacheKey(entry.sid, entry.id); + await fsp.writeFile(path.join(tmp, key), entry.body); + await fsp.writeFile(path.join(tmp, `${key}.json`), JSON.stringify({ readOnly: false })); + } + const tar = spawnSync('tar', ['-czf', '-', '-C', tmp, '.'], { + maxBuffer: 16 * 1024 * 1024, + env: { ...process.env, COPYFILE_DISABLE: '1' }, + }); + await fsp.rm(tmp, { recursive: true, force: true }); + return tar.stdout; +} + +const probe = (refs: unknown) => + fetch(`${baseUrl}/api/v2/session/inputs/probe`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ refs }), + }); + +describe('input delivery routes', () => { + test('disabled shared-runner targets reject input-cache routes without storing bytes', async () => { + config.session_workspace_enabled = false; + const push = await fetch(`${baseUrl}/api/v2/session/inputs`, { + method: 'POST', + headers: { 'Content-Type': 'application/x-gtar' }, + body: 'not-a-tar', + }); + + expect(push.status).toBe(404); + expect((await probe([])).status).toBe(404); + expect(await fsp.lstat(SESSION_INPUT_CACHE_DIR).catch(() => null)).toBeNull(); + }); + + test('probe parses its body and reports everything missing on a cold VM', async () => { + const cacheKey = inputCacheKey('s1', 'f1'); + const response = await probe([{ cache_key: cacheKey }]); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ missing: [{ cache_key: cacheKey }] }); + }); + + test('a pushed batch flips the probe answer to nothing missing', async () => { + const push = await fetch(`${baseUrl}/api/v2/session/inputs`, { + method: 'POST', + headers: { + 'Content-Type': 'application/x-gtar', + 'X-CodeAPI-Input-Expanded-Bytes': String( + Buffer.byteLength('bytes') + + Buffer.byteLength(JSON.stringify({ readOnly: false })), + ), + }, + body: await makeBatch([{ sid: 's1', id: 'f1', body: 'bytes' }]), + }); + expect(push.status).toBe(200); + expect(await push.json()).toEqual({ stored: 1 }); + + const after = await probe([ + { cache_key: inputCacheKey('s1', 'f1') }, + { cache_key: inputCacheKey('s1', 'f2') }, + ]); + /* Only the object the VM actually holds is skipped — dedupe is the VM's + * answer, never control-plane bookkeeping. */ + expect(await after.json()).toEqual({ + missing: [{ cache_key: inputCacheKey('s1', 'f2') }], + }); + }); + + test('probe rejects a malformed ref list rather than guessing', async () => { + expect((await probe('nope')).status).toBe(400); + expect((await probe([{ cache_key: 'not-a-digest' }])).status).toBe(400); + const key = inputCacheKey('s1', 'duplicate'); + expect((await probe([{ cache_key: key }, { cache_key: key }])).status).toBe(400); + expect((await probe( + Array.from( + { length: SESSION_INPUT_CACHE_MAX_OBJECTS + 1 }, + (_, i) => ({ cache_key: inputCacheKey('s1', String(i)) }), + ), + )).status).toBe(400); + }); +}); diff --git a/api/src/api/v2-files.test.ts b/api/src/api/v2-files.test.ts new file mode 100644 index 0000000..540b529 --- /dev/null +++ b/api/src/api/v2-files.test.ts @@ -0,0 +1,88 @@ +import { describe, expect, test } from 'bun:test'; +import { config } from '../config'; +import type { TFile } from '../job'; +import { validateExecuteArguments, validateExecuteFiles } from './v2'; + +function messageOf(fn: () => void): string { + try { + fn(); + } catch (error) { + return (error as { message?: string }).message ?? String(error); + } + return ''; +} + +describe('execute file validation', () => { + test('rejects duplicate and ancestor-conflicting destinations before priming', () => { + expect(messageOf(() => validateExecuteFiles([ + { name: 'data.csv', content: 'a' }, + { name: 'data.csv', content: 'b' }, + ]))).toContain('duplicate destination'); + + expect(messageOf(() => validateExecuteFiles([ + { name: 'results', content: 'file' }, + { name: 'results/out.csv', content: 'nested' }, + ]))).toContain('conflicting destinations'); + }); + + test('rejects malformed stable cache identities', () => { + expect(messageOf(() => validateExecuteFiles([{ + id: 'masked', + storage_session_id: 'masked-session', + name: 'data.csv', + input_cache_key: '../not-a-key', + }]))).toContain('64-character lowercase hex digest'); + }); + + test('rejects ambiguous and type-confused inline/reference shapes', () => { + expect(messageOf(() => validateExecuteFiles([null as unknown as TFile]))) + .toContain('must be an object'); + expect(messageOf(() => validateExecuteFiles([{ + name: 'data.csv', + content: 'inline', + id: 'masked', + storage_session_id: 'masked-session', + }]))).toContain('exactly one'); + expect(messageOf(() => validateExecuteFiles([{ + name: 'data.csv', + content: 'inline', + id: 123, + storage_session_id: {} as string, + input_cache_key: 'a'.repeat(64), + } as unknown as TFile]))).toContain('id must be a non-empty string'); + expect(messageOf(() => validateExecuteFiles([{ + name: 'data.csv', + id: 'masked', + }]))).toContain('storage_session_id'); + expect(messageOf(() => validateExecuteFiles([{ + name: 'data.csv', + content: 'inline', + input_cache_key: 'a'.repeat(64), + }]))).toContain('inline content cannot include'); + }); + + test('validates args and stdin before any workspace priming', () => { + expect(messageOf(() => validateExecuteArguments(123, ''))).toContain('args'); + expect(messageOf(() => validateExecuteArguments(['ok', 123], ''))).toContain('args'); + expect(messageOf(() => validateExecuteArguments([], 123))).toContain('stdin'); + expect(() => validateExecuteArguments(['--flag'], 'input')).not.toThrow(); + }); + + test('caps total destinations even when they all reference one object', () => { + const files = Array.from({ length: config.max_input_files + 1 }, (_, i) => ({ + id: 'masked', + storage_session_id: 'masked-session', + name: `copy-${i}.csv`, + })); + expect(messageOf(() => validateExecuteFiles(files))).toContain('cannot contain more than'); + }); + + test('accepts one object at several independent destinations', () => { + const key = 'a'.repeat(64); + const files: TFile[] = [ + { id: 'masked', storage_session_id: 'masked-session', name: 'a.csv', input_cache_key: key }, + { id: 'masked', storage_session_id: 'masked-session', name: 'copy/a.csv', input_cache_key: key }, + ]; + expect(() => validateExecuteFiles(files)).not.toThrow(); + }); +}); diff --git a/api/src/api/v2-session-binding.test.ts b/api/src/api/v2-session-binding.test.ts new file mode 100644 index 0000000..e5f36a1 --- /dev/null +++ b/api/src/api/v2-session-binding.test.ts @@ -0,0 +1,102 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from 'bun:test'; +import express from 'express'; +import * as fsp from 'fs/promises'; +import type { Server } from 'http'; +import * as os from 'os'; +import * as path from 'path'; +import { config } from '../config'; +import { Job } from '../job'; +import { loadPackage } from '../runtime'; +import { + getBoundSessionWorkspace, + resetSessionWorkspaceStateForTests, + type SessionWorkspace, +} from '../session-workspace'; +import v2Router from './v2'; + +let server: Server; +let baseUrl: string; +let packageDir: string; +const savedSessionWorkspaceEnabled = config.session_workspace_enabled; +const savedRequireExecutionManifest = config.require_execution_manifest; +const testLanguage = 'headerless-session-regression'; +const testVersion = '1.0.0'; + +beforeAll(async () => { + packageDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'headerless-runtime-')); + await fsp.writeFile(path.join(packageDir, 'pkg-info.json'), JSON.stringify({ + language: testLanguage, + version: testVersion, + aliases: [], + })); + loadPackage(packageDir); + + const app = express(); + app.use('/api/v2', v2Router); + await new Promise((resolve) => { + server = app.listen(0, '127.0.0.1', () => resolve()); + }); + const address = server.address(); + baseUrl = `http://127.0.0.1:${typeof address === 'object' && address ? address.port : 0}`; +}); + +afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); + await fsp.rm(packageDir, { recursive: true, force: true }); +}); + +afterEach(() => { + config.session_workspace_enabled = savedSessionWorkspaceEnabled; + config.require_execution_manifest = savedRequireExecutionManifest; + resetSessionWorkspaceStateForTests(); +}); + +const execute = (runtimeSessionId?: string) => + fetch(`${baseUrl}/api/v2/execute`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...(runtimeSessionId ? { 'X-Runtime-Session-Id': runtimeSessionId } : {}), + }, + body: JSON.stringify({ + language: testLanguage, + version: testVersion, + files: [{ name: 'main.txt', content: 'test' }], + }), + }); + +describe('per-request session binding', () => { + test('a headerless execute does not inherit the runner session bound by a prior request', async () => { + config.session_workspace_enabled = true; + config.require_execution_manifest = false; + + const observedSessions: Array = []; + const originalPrime = Job.prototype.prime; + const originalExecute = Job.prototype.execute; + const originalCleanup = Job.prototype.cleanup; + + Job.prototype.prime = async function primeWithoutFilesystem(): Promise { + observedSessions.push( + (this as unknown as { session?: SessionWorkspace }).session, + ); + }; + Job.prototype.execute = async function executeWithoutSandbox() { + return {} as Awaited>; + }; + Job.prototype.cleanup = async function cleanupWithoutFilesystem(): Promise {}; + + try { + expect((await execute('rt_bound_once')).status).toBe(200); + const bound = getBoundSessionWorkspace(); + expect(bound?.runtimeSessionId).toBe('rt_bound_once'); + + expect((await execute()).status).toBe(200); + expect(observedSessions).toEqual([bound, undefined]); + expect(getBoundSessionWorkspace()).toBe(bound); + } finally { + Job.prototype.prime = originalPrime; + Job.prototype.execute = originalExecute; + Job.prototype.cleanup = originalCleanup; + } + }); +}); diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index 5046ca3..36513f6 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -4,7 +4,12 @@ import type { TFile } from '../job'; import { getLatestRuntimeMatchingLanguageVersion, getRuntimes } from '../runtime'; import { logger } from '../logger'; import { config } from '../config'; -import { Job, ValidationError } from '../job'; +import { + Job, + SessionWorkspaceDirtyError, + ValidationError, + validateFilePath, +} from '../job'; import { EXECUTION_MANIFEST_HEADER, ExecutionManifestError, type ExecutionManifestClaims } from '../execution-manifest'; import { verifyExecuteRequestManifest } from '../execution-manifest-request'; import { EGRESS_GRANT_HEADER } from '../egress'; @@ -12,10 +17,112 @@ import { activeSandboxExecutions, recordSandboxExecution } from '../metrics'; import { classifySandboxSafeError } from '../safe-error'; import { withSpan } from '../telemetry'; import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; +import type { SessionWorkspace } from '../session-workspace'; +import { + RUNTIME_SESSION_ID_HEADER, + bindSessionWorkspace, + parseSessionBindingFromHeader, +} from '../session-workspace'; +import { streamSessionCheckpoint, restoreSessionCheckpoint } from '../session-checkpoint'; +import { ensureToolCallSocketProxyReady } from '../tool-call-socket-process'; +import { + SESSION_INPUT_CACHE_MAX_OBJECTS, + hasCachedInput, + pruneInputCache, + storeCachedInputs, +} from '../session-inputs'; const router = express.Router(); const SYNTHETIC_PRINCIPAL_SOURCE = 'synthetic_test'; +function existingDestinationConflictMessage(existing: string, destination: string): string { + return existing === destination + ? `files contains duplicate destination "${destination}"` + : `files contains conflicting destinations "${existing}" and "${destination}"`; +} + +export function validateExecuteFiles(files: TFile[]): void { + if (files.length > config.max_input_files) { + throw { message: `files cannot contain more than ${config.max_input_files} destinations` }; + } + const destinations = new Set(); + for (const [i, value] of files.entries()) { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + throw { message: `files[${i}] must be an object` }; + } + const file = value as TFile; + const inline = typeof file.content === 'string'; + const byRef = typeof file.id === 'string' && file.id.length > 0; + if (inline === byRef) { + throw { + message: `files[${i}] must contain exactly one of non-empty id or string content`, + }; + } + if (file.id !== undefined && !byRef) { + throw { message: `files[${i}].id must be a non-empty string if provided` }; + } + if (byRef) { + if (typeof file.storage_session_id !== 'string' || file.storage_session_id.length === 0) { + throw { message: `files[${i}].storage_session_id is required as a non-empty string for file refs` }; + } + } else if (file.storage_session_id !== undefined || file.input_cache_key !== undefined) { + throw { + message: `files[${i}] inline content cannot include storage_session_id or input_cache_key`, + }; + } + if (file.name !== undefined && typeof file.name !== 'string') { + throw { message: `files[${i}].name must be a string if provided` }; + } + if ( + file.encoding !== undefined + && !(['base64', 'hex', 'utf8'] as const).includes(file.encoding) + ) { + throw { message: `files[${i}].encoding must be base64, hex, or utf8 if provided` }; + } + if (file.entity_id !== undefined && typeof file.entity_id !== 'string') { + throw { message: `files[${i}].entity_id must be a string if provided` }; + } + if ( + file.input_cache_key !== undefined && + ( + typeof file.input_cache_key !== 'string' || + !/^[0-9a-f]{64}$/.test(file.input_cache_key) + ) + ) { + throw { message: `files[${i}].input_cache_key must be a 64-character lowercase hex digest` }; + } + const destination = file.name || `file${i}.code`; + try { + validateFilePath(destination, '/tmp/codeapi-request-validation'); + } catch (error) { + throw { + message: error instanceof Error + ? `files[${i}].name is invalid: ${error.message}` + : `files[${i}].name is invalid`, + }; + } + const conflict = [...destinations].find( + existing => + existing === destination || + existing.startsWith(`${destination}/`) || + destination.startsWith(`${existing}/`), + ); + if (conflict) { + throw { message: existingDestinationConflictMessage(conflict, destination) }; + } + destinations.add(destination); + } +} + +export function validateExecuteArguments(args: unknown, stdin: unknown): void { + if (args !== undefined && (!Array.isArray(args) || args.some(value => typeof value !== 'string'))) { + throw { message: 'args must be an array of strings if provided' }; + } + if (stdin !== undefined && typeof stdin !== 'string') { + throw { message: 'stdin must be a string if provided' }; + } +} + export interface ExecuteRequestBody { /** Top-level execution session id (one sandbox `/exec` invocation). * Intra-monorepo wire — service-api and sandbox ship together, so @@ -130,6 +237,7 @@ function getJob( egressGrantToken?: string, toolCallSocketEnabled = false, isSynthetic = false, + runtimeSessionHeader?: string | string[], ): Job { const { session_id, language, version, args, stdin, files, @@ -151,11 +259,8 @@ function getJob( if (body.tool_call_socket !== undefined && typeof body.tool_call_socket !== 'boolean') { throw { message: 'tool_call_socket must be a boolean if specified' }; } - for (const [i, file] of files.entries()) { - if (typeof file.content !== 'string' && typeof file.id !== 'string') { - throw { message: `files[${i}].content is required as a string if no id is provided` }; - } - } + validateExecuteArguments(args, stdin); + validateExecuteFiles(files); const rt = getLatestRuntimeMatchingLanguageVersion(language, version); if (!rt) { @@ -171,6 +276,31 @@ function getJob( validateConstraints(body, rt); + /* Session mode is per-request opt-in: only run in the persistent workspace + * when THIS request carried a valid X-Runtime-Session-Id. A headerless or + * malformed-header request must NOT inherit a previously bound session, or it + * would reuse that session's files/UID (defense-in-depth — the backend always + * sends the header for a session VM, so this only guards stray requests). */ + const binding = parseSessionBindingFromHeader(runtimeSessionHeader); + let session: SessionWorkspace | null = null; + if (binding) { + session = bindSessionWorkspace(binding) ?? null; + /* The runner is already pinned to a DIFFERENT session. Falling through + * with `session = null` would silently run the request as a stateless + * one-shot under another session's workspace/UID; conflict loudly instead + * so the control plane recycles this VM. */ + if (!session) { + throw { status: 409, message: 'Runner is bound to a different runtime session' }; + } + if (session.dirtyReason) { + throw { + status: 409, + code: 'session_workspace_dirty', + message: 'Session workspace must be restored before another execute', + }; + } + } + return new Job({ session_id: session_id ?? null, runtime: rt, @@ -194,6 +324,7 @@ function getJob( egress_grant: egressGrantToken, tool_call_socket_enabled: toolCallSocketEnabled, is_synthetic: isSynthetic, + session, }); } @@ -238,6 +369,8 @@ function manifestErrorStatus(error: ExecutionManifestError): number { router.use((req: Request, res: Response, next: NextFunction) => { if (['GET', 'HEAD', 'OPTIONS'].includes(req.method)) return next(); + /* Checkpoint restore and additive file delivery stream tar.gz bodies, not JSON. */ + if (req.path === '/session/restore' || req.path === '/session/inputs') return next(); if (!req.headers['content-type']?.startsWith('application/json')) { return res.status(415).json({ message: 'requests must be of type application/json' }); } @@ -259,6 +392,7 @@ router.post('/execute', express.json({ limit: config.execute_body_limit }), asyn let activeExecution = false; let metricsLanguage = 'unknown'; let metricsOutcome: Parameters[0]['outcome'] = 'execution_error'; + let primeCompleted = false; const cleanupHandler = async (): Promise => { if (!job || cleanedUp) return; @@ -322,6 +456,7 @@ router.post('/execute', express.json({ limit: config.execute_body_limit }), asyn tokenFromBodyOrHeader(req.body, 'egress_grant', req.header(EGRESS_GRANT_HEADER) ?? undefined), toolCallSocketEnabled, verifiedManifest?.principal_source === SYNTHETIC_PRINCIPAL_SOURCE, + req.headers[RUNTIME_SESSION_ID_HEADER], ); metricsLanguage = job.runtime.language; markActiveExecution(); @@ -339,13 +474,27 @@ router.post('/execute', express.json({ limit: config.execute_body_limit }), asyn const message = error instanceof Error ? error.message : (error as { message?: unknown })?.message; - return res.status(400).json({ message: message || 'Bad request' }); + /* Most validation paths are 400; a session-binding conflict carries its + * own status so the control plane can tell "bad request" from "this VM + * belongs to another session, recycle me". */ + const status = (error as { status?: unknown })?.status; + const code = (error as { code?: unknown })?.code; + return res + .status(typeof status === 'number' ? status : 400) + .json({ + ...(typeof code === 'string' ? { error: code } : {}), + message: message || 'Bad request', + }); } try { + if (toolCallSocketEnabled) { + await ensureToolCallSocketProxyReady(); + } await withSpan('codeapi.sandbox.prime', { 'codeapi.language': job.runtime.language, }, () => job!.prime()); + primeCompleted = true; const result = await withSpan('codeapi.sandbox.run', { 'codeapi.language': job.runtime.language, }, () => job!.execute()); @@ -387,6 +536,22 @@ router.post('/execute', express.json({ limit: config.execute_body_limit }), asyn metricsOutcome = 'success'; return res.status(200).json(result); } catch (error) { + if (primeCompleted && job?.markSessionDirty('execution failed after input priming')) { + metricsOutcome = 'execution_error'; + logger.error({ job: job.uuid, err: error }, 'Session execution left workspace state unknown'); + return res.status(409).json({ + error: 'session_workspace_dirty', + message: 'Session workspace must be restored before another execute', + }); + } + if (error instanceof SessionWorkspaceDirtyError) { + metricsOutcome = 'execution_error'; + logger.error({ job: job?.uuid, err: error }, 'Session input priming left a partial workspace'); + return res.status(409).json({ + error: error.code, + message: error.message, + }); + } if (error instanceof ValidationError) { metricsOutcome = 'validation_error'; return res.status(400).json({ message: error.message }); @@ -441,4 +606,135 @@ router.get('/runtimes', (_req: Request, res: Response) => { return res.status(200).json(runtimes); }); +/* Session workspace checkpoint / restore — control-plane driven, session-mode + * only. GET streams a tar.gz of the whole workspace (captures state the + * file-ref path drops: installed packages, chDB dirs, unsupported-extension + * files); POST replaces the workspace from one. No body parser on restore: + * the handler consumes the raw request stream. */ +/* Bind the session from the header before checkpoint/restore. These run BEFORE + * the first /execute on a relaunched VM, so in the hookless design nothing else + * has bound the workspace yet; without this the handlers 409 and a real restore + * silently continues with an empty workspace (checkpoint state lost on expiry). + * Returns false (→ fail closed) when this request carries no valid header, so a + * headerless/malformed request never operates on a stale prior session. */ +function bindSessionFromHeader(req: Request): boolean { + const binding = parseSessionBindingFromHeader(req.headers[RUNTIME_SESSION_ID_HEADER]); + if (!binding) return false; + /* A REJECTED bind (this runner is already pinned to a different session) + * must fail the request too — proceeding would run checkpoint/restore/file + * delivery for the requested session against the previously bound session's + * workspace. */ + return bindSessionWorkspace(binding) != null; +} + +/* Express 4 (pinned) does NOT auto-forward rejected route-handler promises, so + * `.catch(next)` is required or a rejection (e.g. session.ownership()) hangs the + * request and surfaces as an unhandled rejection instead of a 5xx. */ +router.get('/session/checkpoint', (req: Request, res: Response, next: NextFunction) => { + if (!bindSessionFromHeader(req)) { + return res.status(409).json({ message: 'Missing runtime session header' }); + } + return streamSessionCheckpoint(res).catch(next); +}); +router.post('/session/restore', (req: Request, res: Response, next: NextFunction) => { + if (!bindSessionFromHeader(req)) { + return res.status(409).json({ message: 'Missing runtime session header' }); + } + return restoreSessionCheckpoint(req, res).catch(next); +}); +/** + * Input delivery for backends whose sandbox cannot reach the file server. + * + * These are deliberately NOT session-scoped: the cache they fill is keyed by + * (storage session, object id) and lives outside any workspace, so the same + * mechanism serves stateful sessions and stateless one-shots alike. The + * workspace is still only ever written by the normal priming path. The routes + * are nevertheless Lambda-runner-only: Lambda's Runtime API supplies the + * external bearer boundary, while ordinary shared runners must not expose an + * unauthenticated cache-write surface. + */ +function requireSessionInputDeliveryTarget( + _req: Request, + res: Response, + next: NextFunction, +): Response | void { + if (!config.session_workspace_enabled) { + return res.status(404).json({ message: 'Not Found' }); + } + next(); +} + +/* No global body parser exists (see index.ts), so the probe installs its own. + * A ref list is tiny — a few hundred bytes per entry — but bound it anyway. */ +router.post( + '/session/inputs/probe', + requireSessionInputDeliveryTarget, + express.json({ limit: '1mb' }), + async (req: Request, res: Response, next: NextFunction) => { + try { + const refs = (req.body as { refs?: Array<{ cache_key?: unknown }> })?.refs; + if (!Array.isArray(refs)) { + return res.status(400).json({ message: 'refs must be an array' }); + } + if (refs.length > SESSION_INPUT_CACHE_MAX_OBJECTS) { + return res.status(400).json({ + message: `refs exceeds the ${SESSION_INPUT_CACHE_MAX_OBJECTS}-object limit`, + }); + } + const missing: Array<{ cache_key: string }> = []; + const seen = new Set(); + for (const ref of refs) { + if (typeof ref?.cache_key !== 'string' || !/^[0-9a-f]{64}$/.test(ref.cache_key)) { + return res.status(400).json({ message: 'each ref requires a valid cache_key' }); + } + if (seen.has(ref.cache_key)) { + return res.status(400).json({ message: 'refs contains a duplicate cache_key' }); + } + seen.add(ref.cache_key); + if (!(await hasCachedInput('', '', ref.cache_key))) { + missing.push({ cache_key: ref.cache_key }); + } + } + return res.status(200).json({ missing }); + } catch (error) { + return next(error); + } + }, +); + +router.post( + '/session/inputs', + requireSessionInputDeliveryTarget, + async (req: Request, res: Response) => { + try { + const expandedHeader = req.get('x-codeapi-input-expanded-bytes'); + let expectedBytes: number | undefined; + if (expandedHeader !== undefined) { + if (!/^(0|[1-9][0-9]*)$/.test(expandedHeader)) { + return res.status(400).json({ message: 'invalid expanded input byte count' }); + } + expectedBytes = Number(expandedHeader); + if ( + !Number.isSafeInteger(expectedBytes) + || expectedBytes > config.input_cache_max_bytes + ) { + return res.status(400).json({ message: 'invalid expanded input byte count' }); + } + } + const stored = await storeCachedInputs( + req, + config.input_cache_max_bytes, + expectedBytes, + ); + await pruneInputCache(config.input_cache_max_bytes).catch((err) => { + logger.warn({ err }, 'Failed to prune session input cache'); + }); + return res.status(200).json({ stored }); + } catch (error) { + logger.error({ err: error }, 'Failed to store session inputs'); + return res.status(500).json({ message: 'session input delivery failed' }); + } + }, +); + export default router; diff --git a/api/src/config.ts b/api/src/config.ts index f8e2aba..7724ff0 100644 --- a/api/src/config.ts +++ b/api/src/config.ts @@ -65,6 +65,12 @@ export const config = { run_memory_limit: Number(process.env.SANDBOX_RUN_MEMORY_LIMIT ?? -1), max_concurrent_jobs: safeInt(process.env.SANDBOX_MAX_CONCURRENT_JOBS, 8), per_job_uids: (process.env.SANDBOX_PER_JOB_UIDS ?? 'true') === 'true', + /* Image-level enable for persistent session workspaces (stateful sessions). + * Only the Lambda MicroVM runner target sets this true; the K8s + * sandbox-runner image leaves it false so it is structurally incapable of + * session mode. An enabled runner additionally binds each request to a + * workspace through the authenticated X-Runtime-Session-Id header. */ + session_workspace_enabled: (process.env.SANDBOX_SESSION_WORKSPACE_ENABLED ?? 'false') === 'true', job_uid_base: safeInt(process.env.SANDBOX_JOB_UID_BASE, 200000), job_gid_base: safeInt(process.env.SANDBOX_JOB_GID_BASE, 200000), job_uid_count: safeInt( @@ -77,6 +83,25 @@ export const config = { nsjail_path: process.env.NSJAIL_PATH ?? '/usr/sbin/nsjail', nsjail_config: process.env.NSJAIL_CONFIG ?? '/sandbox_api/config/sandbox.cfg', execute_body_limit: process.env.SANDBOX_EXECUTE_BODY_LIMIT ?? '50mb', + /* Independent runner-side ceiling for streamed checkpoint restores. The + * control plane enforces the same default, but the receiver must not trust a + * caller-supplied Content-Length or upstream enforcement. */ + checkpoint_max_bytes: safeInt( + process.env.SANDBOX_CHECKPOINT_MAX_BYTES, + 512 * 1024 * 1024, + ), + /* Ceiling for the pushed input cache (session-inputs.ts). Eviction is + * always safe — a miss simply re-pushes on the next probe — so this is a + * disk guard, not a correctness knob. */ + input_cache_max_bytes: safeInt( + process.env.SANDBOX_INPUT_CACHE_MAX_BYTES, + 512 * 1024 * 1024, + ), + /* Bound both validation cost and per-request priming fan-out. The cache's + * unique-object cap is not enough because one object may be requested at + * many destinations. */ + max_input_files: safeInt(process.env.SANDBOX_MAX_INPUT_FILES, 256), + prime_concurrency: safeInt(process.env.SANDBOX_PRIME_CONCURRENCY, 8), egress_gateway_url: egressGatewayUrl, file_server_url: process.env.FILE_SERVER_URL ?? '', max_nesting_depth: safeInt(process.env.SANDBOX_MAX_NESTING_DEPTH, 10), diff --git a/api/src/download.test.ts b/api/src/download.test.ts index d9f20c5..692957b 100644 --- a/api/src/download.test.ts +++ b/api/src/download.test.ts @@ -3,11 +3,16 @@ import * as fsp from 'fs/promises'; import * as path from 'path'; import * as os from 'os'; import * as semver from 'semver'; -import { Job, type TFile } from './job'; +import { Job, SessionWorkspaceDirtyError, type TFile } from './job'; import type { Runtime } from './runtime'; import { config } from './config'; import { SANDBOX_DIR_MODE, SANDBOX_FILE_MODE } from './validation'; -import { SANDBOX_READONLY_FILE_MODE, compatibilityModeForSkippedChown } from './workspace-isolation'; +import type { SessionWorkspace } from './session-workspace'; +import { + SANDBOX_READONLY_FILE_MODE, + compatibilityModeForSkippedChown, + fallbackSandboxIdentity, +} from './workspace-isolation'; /** * Integration tests for `Job.downloadAndWriteFile` against a real HTTP @@ -47,7 +52,7 @@ function makeRuntime(): Runtime { }; } -function makeJob(files: TFile[] = []): Job { +function makeJob(files: TFile[] = [], session?: SessionWorkspace): Job { return new Job({ session_id: 'test-session', runtime: makeRuntime(), @@ -57,9 +62,29 @@ function makeJob(files: TFile[] = []): Job { timeouts: { compile: 5000, run: 5000 }, cpu_times: { compile: 5000, run: 5000 }, memory_limits: { compile: 100_000_000, run: 100_000_000 }, + session, }); } +function sessionWorkspaceAt( + dir: string, + runtimeSessionId: string, + markDirty: () => void = () => {}, +): SessionWorkspace { + const identity = fallbackSandboxIdentity(); + return { + runtimeSessionId, + acquire: async () => ({ + workspaceId: runtimeSessionId, + dir, + identity, + }), + primedInputId: () => undefined, + markPrimed: () => {}, + markDirty, + } as unknown as SessionWorkspace; +} + function currentUid(): number | undefined { return typeof process.getuid === 'function' ? process.getuid() : undefined; } @@ -76,6 +101,7 @@ type Route = { contentDisposition?: string; headers?: Record; body?: string; + delayMs?: number; onRequest?: (req: Request) => void; }; @@ -92,11 +118,14 @@ beforeAll(() => { originalPerJobUids = config.per_job_uids; server = Bun.serve({ port: 0, - fetch(req) { + async fetch(req) { const url = new URL(req.url); const route = routes.get(url.pathname); if (!route) return new Response('not found', { status: 404 }); route.onRequest?.(req); + if (route.delayMs) { + await new Promise(resolve => setTimeout(resolve, route.delayMs)); + } const headers = new Headers(); if (route.contentDisposition) { headers.set('content-disposition', route.contentDisposition); @@ -260,6 +289,103 @@ describe('downloadAndWriteFile / RFC 5987 round-trip', () => { expect(contents).toBe('legacy bytes'); }); + it('resolves concurrent header destinations without provisional-name false conflicts', async () => { + const renamed: TFile = { + id: 'renamed-id', + storage_session_id: 'prev-session', + name: 'vacated.txt', + }; + const replacement: TFile = { + id: 'replacement-id', + storage_session_id: 'prev-session', + name: 'actual.txt', + }; + routes.set(`/sessions/${encodeURIComponent(renamed.storage_session_id!)}/objects/${encodeURIComponent(renamed.id!)}`, { + status: 200, + contentDisposition: 'attachment; filename="actual.txt"', + body: 'renamed bytes', + /* Make the other ref resolve `vacated.txt` while this ref's requested + * name would still be provisional under the old reservation scheme. */ + delayMs: 75, + }); + routes.set(`/sessions/${encodeURIComponent(replacement.storage_session_id!)}/objects/${encodeURIComponent(replacement.id!)}`, { + status: 200, + contentDisposition: 'attachment; filename="vacated.txt"', + body: 'replacement bytes', + }); + + const session = sessionWorkspaceAt(tmpDir, 'rt_concurrent_rename'); + const job = makeJob([renamed, replacement], session); + const originalPrimeConcurrency = config.prime_concurrency; + config.prime_concurrency = 2; + try { + await job.prime(); + const submissionDir = asInternals(job).submissionDir; + expect(await fsp.readFile(path.join(submissionDir, 'actual.txt'), 'utf8')) + .toBe('renamed bytes'); + expect(await fsp.readFile(path.join(submissionDir, 'vacated.txt'), 'utf8')) + .toBe('replacement bytes'); + } finally { + config.prime_concurrency = originalPrimeConcurrency; + await job.cleanup(); + } + }); + + it('rejects concurrent refs that resolve to the same destination before either can overwrite', async () => { + const slower: TFile = { + id: 'same-slower-id', + storage_session_id: 'prev-session', + name: 'slower-fallback.txt', + }; + const faster: TFile = { + id: 'same-faster-id', + storage_session_id: 'prev-session', + name: 'faster-fallback.txt', + }; + routes.set(`/sessions/${encodeURIComponent(slower.storage_session_id!)}/objects/${encodeURIComponent(slower.id!)}`, { + status: 200, + contentDisposition: 'attachment; filename="same.txt"', + body: 'slower bytes', + delayMs: 75, + }); + routes.set(`/sessions/${encodeURIComponent(faster.storage_session_id!)}/objects/${encodeURIComponent(faster.id!)}`, { + status: 200, + contentDisposition: 'attachment; filename="same.txt"', + body: 'faster bytes', + }); + + let dirty = false; + const job = makeJob( + [slower, faster], + sessionWorkspaceAt(tmpDir, 'rt_concurrent_same_destination', () => { dirty = true; }), + ); + const originalPrimeConcurrency = config.prime_concurrency; + config.prime_concurrency = 2; + let deadlockTimer: ReturnType | undefined; + try { + const outcome = await Promise.race([ + job.prime().then( + () => ({ status: 'fulfilled' as const }), + error => ({ status: 'rejected' as const, error }), + ), + new Promise<{ status: 'timeout' }>(resolve => { + deadlockTimer = setTimeout(() => resolve({ status: 'timeout' }), 2_000); + }), + ]); + if (deadlockTimer) clearTimeout(deadlockTimer); + expect(outcome.status).toBe('rejected'); + if (outcome.status === 'rejected') { + expect(outcome.error).toBeInstanceOf(SessionWorkspaceDirtyError); + } + expect(dirty).toBe(true); + expect(await fsp.readFile(path.join(tmpDir, 'same.txt'), 'utf8')).toBe('faster bytes'); + } finally { + if (deadlockTimer) clearTimeout(deadlockTimer); + config.prime_concurrency = originalPrimeConcurrency; + await job.cleanup(); + } + }); + it('decodes UTF-8 percent-encoded names with non-ASCII characters', async () => { const file: TFile = { id: 'utf8-id', @@ -282,7 +408,7 @@ describe('downloadAndWriteFile / RFC 5987 round-trip', () => { expect(contents).toBe('hi'); }); - it('returns null when the server keeps 404-ing past the retry cap (no phantom write)', async () => { + it('fails when the server keeps 404-ing past the retry cap (no phantom write)', async () => { const file: TFile = { id: 'missing-id', storage_session_id: 'prev-session', @@ -293,9 +419,7 @@ describe('downloadAndWriteFile / RFC 5987 round-trip', () => { const job = makeJob([file]); asInternals(job).submissionDir = tmpDir; - const writtenName = await job.downloadAndWriteFile(file, 2, 1); - - expect(writtenName).toBeNull(); + await expect(job.downloadAndWriteFile(file, 2, 1)).rejects.toThrow('HTTP error: 404'); /* Defensive: confirm we did not leave a partial / phantom file on * disk after exhausting retries. */ await expect(fsp.access(path.join(tmpDir, 'should-not-exist.txt'))).rejects.toThrow(); diff --git a/api/src/entrypoint.sh b/api/src/entrypoint.sh index 69145b1..b532a67 100755 --- a/api/src/entrypoint.sh +++ b/api/src/entrypoint.sh @@ -3,6 +3,23 @@ set -e echo "Starting NsJail sandbox API..." +# Raise the RLIMIT_NOFILE limits so per-job nsjail can set its own soft +# limit (SANDBOX_MAX_OPEN_FILES) without EPERM. The AWS Lambda MicroVM base +# image ships a 1024 hard cap, below the sandbox default of 2048; nsjail +# runs the child with keep_caps:false, so the hard limit must already be +# high before the API starts. Only ever RAISE, never lower: `ulimit -n` +# clamps downward too, so an unconditional set would shrink Docker's ~1M +# default (and any SANDBOX_MAX_OPEN_FILES above 65536) instead of no-op'ing. +raise_nofile() { + local kind="$1" want=65536 cur + cur="$(ulimit "$kind" 2>/dev/null || echo unlimited)" + if [ "$cur" != "unlimited" ] && [ "$cur" -lt "$want" ] 2>/dev/null; then + ulimit "$kind" "$want" 2>/dev/null || true + fi +} +raise_nofile -Hn # hard first, so the soft raise below is permitted +raise_nofile -Sn + SANDBOX_USE_CGROUPV2="${SANDBOX_USE_CGROUPV2:-true}" SANDBOX_REMOVE_UMOUNT_AFTER_STARTUP="${SANDBOX_REMOVE_UMOUNT_AFTER_STARTUP:-true}" NSJAIL_CONFIG_SOURCE="${NSJAIL_CONFIG:-/sandbox_api/config/sandbox.cfg}" @@ -119,24 +136,11 @@ if [ "$ALLOWED_PORT" -gt 0 ] 2>/dev/null; then exit 1 fi - echo "Configuring tool call server forwarding for UID $SANDBOX_UID (port $ALLOWED_PORT)" - - # Start a narrow Unix-socket proxy for sandbox-originated tool calls. - # NsJail bind-mounts this socket and runs a relay inside the sandbox, - # keeping clone_newnet: true (fully isolated network namespace). - # Only POST /tool-call is exposed; health, readiness, metrics, and - # internal gateway routes remain outside the sandbox contract. - FORWARD_TARGET="${SANDBOX_FORWARD_TARGET:-}" - TCS_SOCKET="/tmp/tcs.sock" - if [ -n "$FORWARD_TARGET" ]; then - echo "Starting tool-call socket proxy: $TCS_SOCKET -> $FORWARD_TARGET/tool-call" - # Run under Node, not Bun: Bun's node:http compat layer never fires - # 'connection' events and Bun.serve's idleTimeout does not close - # silent unix-socket connections, which silently disables the - # proxy's DoS defenses. The .build artifact is produced at image - # build time by `bun build --target=node`. See api/Dockerfile. - TCS_SOCKET="$TCS_SOCKET" TCS_SOCKET_UID="$SANDBOX_UID" TCS_SOCKET_GID="$SANDBOX_UID" SANDBOX_FORWARD_TARGET="$FORWARD_TARGET" node /sandbox_api/.build/tool-call-socket-proxy.cjs & - fi + echo "Tool-call socket forwarding enabled for UID $SANDBOX_UID (port $ALLOWED_PORT)" + # Do not start Node here. Lambda MicroVM image creation snapshots this + # already-running container, and the official Node binary embeds OpenSSL. + # The API starts and awaits the narrow Unix-socket proxy only after a + # post-restore /execute is authorized for tool calls. fi # Package permissions are finalized by package-init when the PVC is populated. diff --git a/api/src/execution-manifest-request.test.ts b/api/src/execution-manifest-request.test.ts index bcb9c97..4e7c92b 100644 --- a/api/src/execution-manifest-request.test.ts +++ b/api/src/execution-manifest-request.test.ts @@ -249,6 +249,23 @@ describe('execute request manifest validation', () => { })).toEqual(legacyClaims); }); + test('never admits input cache selectors through the legacy no-body-hash grace', () => { + const cachedBody = { + ...body, + files: [ + body.files[0], + { ...body.files[1], input_cache_key: 'a'.repeat(64) }, + ], + }; + expectManifestError(() => verifyExecuteRequestManifest({ + headerValue: signExecutionManifest(claims(), SECRET), + secret: SECRET, + body: cachedBody, + nowSeconds: 150, + bodyHashRequiredAfterSeconds: 200, + }), 'scope_mismatch'); + }); + test('rejects scoped legacy manifests after the body-hash rollout grace window', () => { expectManifestError(() => verifyExecuteRequestManifest({ headerValue: signExecutionManifest(claims(), SECRET), diff --git a/api/src/execution-manifest-request.ts b/api/src/execution-manifest-request.ts index 6684bef..08b65b0 100644 --- a/api/src/execution-manifest-request.ts +++ b/api/src/execution-manifest-request.ts @@ -48,6 +48,21 @@ function assertManifestBodyHashMatches( options: ManifestBodyHashOptions = {}, ): void { if (!manifest.execute_body_sha256) { + /* input_cache_key selects runner-local bytes and is deliberately absent + * from the coarse input-file claim tuple. It is therefore safe only when + * the whole execute body is signed; never admit it through the rolling + * legacy no-body-hash grace. */ + const hasInputCacheSelector = Array.isArray(body.files) && body.files.some( + file => file != null + && typeof file === 'object' + && typeof (file as TFile).input_cache_key === 'string', + ); + if (hasInputCacheSelector) { + throw new ExecutionManifestError( + 'scope_mismatch', + 'Execution manifest body hash is required for input cache selectors', + ); + } const nowSeconds = options.nowSeconds ?? Math.floor(Date.now() / 1000); /* Updated sandbox pods can receive still-valid manifests from older * service workers during a rolling deploy. Keep that compatibility window diff --git a/api/src/index.ts b/api/src/index.ts index 92c590d..629436e 100644 --- a/api/src/index.ts +++ b/api/src/index.ts @@ -6,7 +6,10 @@ import { validateHardenedSandboxStartup } from './secure-startup'; import { initializeSandboxWorkspaceIsolation, startWorkspaceReaper } from './workspace-isolation'; import { httpMetricsMiddleware, metricsHandler } from './metrics'; import { positiveInt, shutdownTelemetry, traceHttpRequest } from './telemetry'; +import { startWarmupCommand } from './warmup'; +import { stopToolCallSocketProxy } from './tool-call-socket-process'; import v2Router from './api/v2'; +import lifecycleRouter, { LIFECYCLE_HOOK_BASE_PATH } from './api/lifecycle'; const app = express(); @@ -29,6 +32,7 @@ loadPackages(config.packages_directory); logger.info('Registering routes'); app.get('/metrics', metricsHandler); +app.use(LIFECYCLE_HOOK_BASE_PATH, lifecycleRouter); app.use('/api/v2', v2Router); app.get('/', (_req, res) => { @@ -64,13 +68,13 @@ app.use((err: HttpError, _req: express.Request, res: express.Response, _next: ex async function main(): Promise { validateHardenedSandboxStartup(); await initializeSandboxWorkspaceIsolation(); + await startWarmupCommand(); const [address, port] = config.bind_address.split(':'); const stopWorkspaceReaper = startWorkspaceReaper(); const server = app.listen(Number(port), address, () => { logger.info({ address: config.bind_address }, 'Sandbox API started'); }); - let shuttingDown = false; const closeHttpServer = (): Promise => new Promise((resolve, reject) => { server.close((error) => { @@ -107,6 +111,9 @@ async function main(): Promise { shuttingDown = true; stopWorkspaceReaper(); await closeHttpServerWithTimeout(); + await stopToolCallSocketProxy().catch((err) => { + logger.warn({ err }, 'Tool-call socket proxy shutdown failed'); + }); try { await shutdownTelemetry(); } catch (err) { diff --git a/api/src/job-cleanup.test.ts b/api/src/job-cleanup.test.ts index 88f8785..de4b52e 100644 --- a/api/src/job-cleanup.test.ts +++ b/api/src/job-cleanup.test.ts @@ -1,8 +1,22 @@ import { describe, expect, test } from 'bun:test'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; import * as semver from 'semver'; -import { Job } from './job'; +import { + Job, + SessionWorkspaceDirtyError, + ValidationError, + type TFile, +} from './job'; import type { Runtime } from './runtime'; -import { sandboxJobUidPool, type SandboxJobIdentity } from './workspace-isolation'; +import type { SessionWorkspace } from './session-workspace'; +import { + sandboxJobUidPool, + type SandboxJobIdentity, + type SandboxWorkspaceLease, +} from './workspace-isolation'; +import { config } from './config'; interface CleanupInternals { jobIdentity?: SandboxJobIdentity; @@ -59,4 +73,127 @@ describe('Job cleanup', () => { expect(sandboxJobUidPool.availableCount()).toBe(availableBefore); expect(sandboxJobUidPool.activeCount()).toBe(activeBefore); }); + + test('prime waits for every sibling operation before exposing a session failure', async () => { + const workspace = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-settlement-')); + const identity: SandboxJobIdentity = { + slot: 0, + uid: typeof process.getuid === 'function' ? process.getuid() : 0, + gid: typeof process.getgid === 'function' ? process.getgid() : 0, + perJobUid: false, + }; + const lease: SandboxWorkspaceLease = { + workspaceId: 'prime-settlement', + dir: workspace, + identity, + }; + let dirty = false; + const session = { + runtimeSessionId: 'rt_prime_settlement', + acquire: async () => lease, + markDirty: () => { + dirty = true; + }, + } as unknown as SessionWorkspace; + const files: TFile[] = [ + { name: 'fast-failure.txt', content: 'fast' }, + { name: 'slow-sibling.txt', content: 'slow' }, + ]; + const job = new Job({ + session_id: 'prime-settlement', + runtime: makeRuntime(), + files, + args: [], + stdin: '', + timeouts: { compile: 5000, run: 5000 }, + cpu_times: { compile: 5000, run: 5000 }, + memory_limits: { compile: 100_000_000, run: 100_000_000 }, + session, + }); + type PrimeContext = { + submissionDir: string; + identity: SandboxJobIdentity; + signal?: AbortSignal; + }; + const internals = job as unknown as { + writeFile(file: TFile, context?: PrimeContext): Promise; + }; + let slowSettled = false; + internals.writeFile = async (file, context) => { + if (file.name === 'fast-failure.txt') { + throw new ValidationError('scripted prime failure'); + } + await new Promise(resolve => setTimeout(resolve, 30)); + /* Ignore the sibling abort on purpose. The aggregate still must wait, + * and the immutable context must keep this write inside the workspace. */ + await fsp.writeFile(path.join(context!.submissionDir, file.name), file.content!); + slowSettled = true; + }; + + try { + await expect(job.prime()).rejects.toBeInstanceOf(SessionWorkspaceDirtyError); + expect(slowSettled).toBe(true); + expect(dirty).toBe(true); + expect(await fsp.readFile(path.join(workspace, 'slow-sibling.txt'), 'utf8')).toBe('slow'); + await job.cleanup(); + expect(await fsp.lstat(path.join(process.cwd(), 'slow-sibling.txt')).catch(() => null)).toBeNull(); + } finally { + await fsp.rm(workspace, { recursive: true, force: true }); + } + }); + + test('bounds concurrent priming operations', async () => { + const workspace = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-concurrency-')); + const identity: SandboxJobIdentity = { + slot: 0, + uid: typeof process.getuid === 'function' ? process.getuid() : 0, + gid: typeof process.getgid === 'function' ? process.getgid() : 0, + perJobUid: false, + }; + const lease: SandboxWorkspaceLease = { + workspaceId: 'prime-concurrency', + dir: workspace, + identity, + }; + const session = { + runtimeSessionId: 'rt_prime_concurrency', + acquire: async () => lease, + markDirty: () => {}, + } as unknown as SessionWorkspace; + const files: TFile[] = Array.from( + { length: config.prime_concurrency + 4 }, + (_, i) => ({ name: `input-${i}.txt`, content: String(i) }), + ); + const job = new Job({ + session_id: 'prime-concurrency', + runtime: makeRuntime(), + files, + args: [], + stdin: '', + timeouts: { compile: 5000, run: 5000 }, + cpu_times: { compile: 5000, run: 5000 }, + memory_limits: { compile: 100_000_000, run: 100_000_000 }, + session, + }); + const internals = job as unknown as { + writeFile(file: TFile): Promise; + }; + let active = 0; + let maxActive = 0; + internals.writeFile = async () => { + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise(resolve => setTimeout(resolve, 10)); + active -= 1; + }; + + try { + await job.prime(); + expect(maxActive).toBeLessThanOrEqual(config.prime_concurrency); + expect(maxActive).toBeGreaterThan(1); + } finally { + await job.cleanup(); + await fsp.rm(workspace, { recursive: true, force: true }); + } + }); }); diff --git a/api/src/job.ts b/api/src/job.ts index 291d76d..4a57971 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -22,12 +22,14 @@ import { applySandboxPathPermissionsNoFollow, cleanupSandboxWorkspace, createSandboxWorkspace, + ensureDirNoFollow, fallbackSandboxIdentity, retainWorkspaceCleanupUntilRemoved, sandboxJobUidPool, type SandboxJobIdentity, type SandboxWorkspaceLease, } from './workspace-isolation'; +import type { SessionWorkspace } from './session-workspace'; import { DIRKEEP, SANDBOX_DIR_MODE, @@ -38,6 +40,7 @@ import { validateFilePath, isValidFilePath, } from './validation'; +import { cachedInputResponse, inputCacheKey, openCachedInput } from './session-inputs'; export { DIRKEEP, @@ -573,6 +576,9 @@ export interface TFile { * storage). Distinct from the top-level execution session of a `/exec` * call — those are different concepts and were historically conflated. */ storage_session_id?: string; + /** Stable opaque identity supplied by the control plane when sandbox-visible + * ids are per-execution handles. Never used as a filesystem path directly. */ + input_cache_key?: string; name: string; content?: string; encoding?: 'base64' | 'hex' | 'utf8'; @@ -628,6 +634,24 @@ interface InputFileInfo { readOnly?: boolean; } +interface PrimeOperationContext { + /** Captured before any file operation starts. Cleanup clears the Job fields, + * so asynchronous priming code must never re-read them after an await. */ + submissionDir: string; + identity: SandboxJobIdentity; + signal?: AbortSignal; +} + +export class SessionWorkspaceDirtyError extends Error { + readonly code = 'session_workspace_dirty'; + + constructor(cause?: unknown) { + super('Session input priming did not complete; the workspace must be restored'); + this.name = 'SessionWorkspaceDirtyError'; + this.cause = cause; + } +} + interface ExecuteResult { compile?: NsJailResult; run?: NsJailResult; @@ -675,11 +699,21 @@ export class Job { private workspaceLease: SandboxWorkspaceLease | undefined; private jobIdentity: SandboxJobIdentity | undefined; private generatedFiles: GeneratedFile[] = []; + /** Session output-diffing: fileId → {name, signature} pending a surfaced-mark. + * Recorded during the workspace scan but only committed to `session.surfaced` + * AFTER the file uploads, so a dropped upload doesn't permanently suppress an + * unchanged file on the next turn. */ + private pendingSurfaced = new Map(); private sessionFiles: FileRef[] = []; private inheritedRefs: FileRef[] = []; private inputFileHashes = new Map(); + private inputDestinations = new Map(); private entryPointName: string | undefined; private chmoddedDirs = new Set(); + /* Persistent session workspace (stateful mode). When set, the job reuses + * one long-lived workspace + pinned UID across calls instead of a fresh + * per-job workspace; undefined = legacy fresh-per-job path (unchanged). */ + private readonly session: SessionWorkspace | undefined; constructor(opts: { /** Top-level execution session id. Becomes `Job.uuid` and is the id @@ -698,7 +732,10 @@ export class Job { egress_grant?: string; tool_call_socket_enabled?: boolean; is_synthetic?: boolean; + /* Injected when this VM is bound to a stateful runtime session. */ + session?: SessionWorkspace | null; }) { + this.session = opts.session ?? undefined; this.uuid = opts.session_id ?? nanoid(); this.outputSessionId = opts.output_session_id ?? this.uuid; this.log = rootLogger.child({ job: this.uuid }); @@ -710,6 +747,7 @@ export class Job { * historically these collapsed onto the same `session_id` field * which is exactly the conflation this rename eliminates. */ storage_session_id: file.storage_session_id ?? this.outputSessionId, + input_cache_key: file.input_cache_key, name: file.name || `file${i}.code`, content: file.content, encoding: (['base64', 'hex', 'utf8'] as const).includes(file.encoding as 'base64' | 'hex' | 'utf8') @@ -733,13 +771,29 @@ export class Job { this.isSynthetic = opts.is_synthetic === true; } + /** Marks a persistent workspace unusable after a post-prime failure. Returns + * false for stateless jobs so the route can retain its ordinary error shape. */ + markSessionDirty(reason: string): boolean { + if (!this.session) return false; + this.session.markDirty(reason); + return true; + } + async computeFileHash(filePath: string, noFollow = false): Promise { const hash = crypto.createHash('sha256'); - const stream = fs.createReadStream(filePath, noFollow - ? { flags: fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW } - : undefined); - for await (const chunk of stream) hash.update(chunk as Buffer); - return hash.digest('hex'); + /* Numeric open flags are only typed on fsp.open; `createReadStream`'s + * options want a string mode with no O_NOFOLLOW spelling, so open the + * handle first and stream from it. */ + const handle = noFollow + ? await fsp.open(filePath, fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW) + : undefined; + try { + const stream = handle ? handle.createReadStream() : fs.createReadStream(filePath); + for await (const chunk of stream) hash.update(chunk as Buffer); + return hash.digest('hex'); + } finally { + await handle?.close().catch(() => {}); + } } /** @@ -762,12 +816,16 @@ export class Job { return this.jobIdentity ?? fallbackSandboxIdentity(); } - private async applySandboxFilePermissions(filePath: string, noFollow = false): Promise { + private async applySandboxFilePermissions( + filePath: string, + noFollow = false, + identity = this.sandboxIdentity(), + ): Promise { if (noFollow) { - await applySandboxPathPermissionsNoFollow(filePath, this.sandboxIdentity(), SANDBOX_FILE_MODE, 'file'); + await applySandboxPathPermissionsNoFollow(filePath, identity, SANDBOX_FILE_MODE, 'file'); return; } - await applySandboxPathPermissions(filePath, this.sandboxIdentity(), SANDBOX_FILE_MODE); + await applySandboxPathPermissions(filePath, identity, SANDBOX_FILE_MODE); } /** @@ -775,25 +833,118 @@ export class Job { * (inclusive) so the per-job outside UID can create siblings/children while * escaped sibling UIDs cannot traverse the workspace tree. */ - private async secureAncestors(leaf: string): Promise { - const rel = path.relative(this.submissionDir, leaf); + /** + * Creates every directory from submissionDir (exclusive) to `target` + * (inclusive) one component at a time, refusing to descend through a symlink. + * `fsp.mkdir(dir, { recursive: true })` would instead follow a symlinked + * ancestor a prior session turn planted and create directories OUTSIDE the + * workspace as root — so a persistent-session prime must build the path with + * no-follow semantics before writing. + */ + private ensureDirNoFollow( + target: string, + submissionDir = this.submissionDir, + ): Promise { + /* Shared with the pushed-delivery merge (session-checkpoint.ts) so both + * privileged writers into a persistent workspace enforce identical + * no-follow semantics. `secureAncestors` applies the ownership pass here, + * so no identity is handed to the shared helper. */ + return ensureDirNoFollow(submissionDir, target); + } + + private async secureAncestors( + leaf: string, + submissionDir = this.submissionDir, + identity = this.sandboxIdentity(), + ): Promise { + const rel = path.relative(submissionDir, leaf); if (!rel || rel === '..' || rel.startsWith('..' + path.sep)) return; const parts = rel.split(path.sep).filter(Boolean); - let cursor = this.submissionDir; + let cursor = submissionDir; for (const part of parts) { cursor = path.join(cursor, part); /* Parallel downloads under shared parent dirs call into this method * concurrently. Skip paths we've already chmodded to avoid N*M redundant * syscalls (N files × M shared ancestors). */ if (this.chmoddedDirs.has(cursor)) continue; - await applySandboxPathPermissions(cursor, this.sandboxIdentity(), SANDBOX_DIR_MODE); + /* A persistent session workspace can contain a symlink planted by a + * prior turn's sandbox code; priming through it as root would follow the + * link and write/chmod an arbitrary target outside the workspace. Reject + * a symlinked ancestor (lstat never follows) before touching it. Fresh + * per-job workspaces never contain one, so this is a no-op there. */ + const st = await fsp.lstat(cursor); + if (st.isSymbolicLink()) { + throw new Error(`Refusing to prime through symlinked workspace path: ${path.relative(submissionDir, cursor)}`); + } + await applySandboxPathPermissions(cursor, identity, SANDBOX_DIR_MODE); this.chmoddedDirs.add(cursor); } } + /** + * Atomically replaces this ref's provisional/previous aliases with its + * authoritative destination. There is no await in this transition: two + * concurrent downloads resolving to the same path cannot both pass the + * conflict check and start writing. + */ + private reserveInputDestination(file: TFile, destination: string): void { + const conflict = [...this.inputDestinations.entries()].find( + ([reserved, owner]) => + owner !== file + && ( + reserved === destination + || reserved.startsWith(`${destination}/`) + || destination.startsWith(`${reserved}/`) + ), + ); + if (conflict) { + throw new ValidationError( + `Conflicting input destinations: ${conflict[0]} and ${destination}`, + ); + } + for (const [reserved, owner] of this.inputDestinations) { + if (owner === file && reserved !== destination) { + this.inputDestinations.delete(reserved); + } + } + this.inputDestinations.set(destination, file); + } + async prime(): Promise { - this.jobIdentity = await acquireJobIdentity(this.log); - this.workspaceLease = await createSandboxWorkspace(this.jobIdentity); + this.inputDestinations.clear(); + const requestedDestinations = new Map(); + for (const file of this.files) { + validateFilePath(file.name, '/tmp/codeapi-request-validation'); + const conflict = [...requestedDestinations.entries()].find( + ([destination, owner]) => + owner !== file && + ( + destination === file.name || + destination.startsWith(`${file.name}/`) || + file.name.startsWith(`${destination}/`) + ), + ); + if (conflict) { + throw new ValidationError( + `Conflicting input destinations: ${conflict[0]} and ${file.name}`, + ); + } + requestedDestinations.set(file.name, file); + /* Inline destinations are final, so keep them reserved while reference + * downloads resolve their authoritative Content-Disposition names. + * A ref's requested name is only a fallback, not a real destination yet: + * reserving every ref here makes concurrent swaps/order-dependent + * renames falsely conflict before the owning response has resolved. */ + if (!file.id) this.inputDestinations.set(file.name, file); + } + + if (this.session) { + this.workspaceLease = await this.session.acquire(); + this.jobIdentity = this.workspaceLease.identity; + } else { + this.jobIdentity = await acquireJobIdentity(this.log); + this.workspaceLease = await createSandboxWorkspace(this.jobIdentity); + } this.submissionDir = this.workspaceLease.dir; if (!this.isSynthetic) { @@ -803,6 +954,7 @@ export class Job { workspaceId: this.workspaceLease.workspaceId, uid: this.jobIdentity.uid, gid: this.jobIdentity.gid, + session: this.session ? this.session.runtimeSessionId : undefined, }, 'Priming job', ); @@ -812,15 +964,142 @@ export class Job { await this.autoLoadDirkeep(); } - const fileOps: Promise[] = []; + /* Promise.all rejects as soon as one operation fails, while its siblings + * keep running. The route's finally then calls cleanup(), which clears the + * session path/identity. A delayed sibling used to resume afterward and + * derive a relative destination from `submissionDir === ''`, writing under + * the runner's own cwd. Abort the siblings and, critically, wait for every + * operation to settle before exposing the failure to request cleanup. */ + const controller = new AbortController(); + const operationBase = { + submissionDir: this.submissionDir, + identity: this.jobIdentity, + }; + let firstFailure: { error: unknown } | undefined; + const runFileOperation = async (operation: () => Promise): Promise => { + try { + await operation(); + } catch (error) { + if (!firstFailure) { + firstFailure = { error }; + controller.abort(error); + } + throw error; + } + }; + + const fileOps: Array<() => Promise> = []; for (const file of this.files) { + const operationContext: PrimeOperationContext = { + ...operationBase, + signal: controller.signal, + }; if (file.id) { - fileOps.push(this.downloadAndWriteFile(file).then(() => {})); + fileOps.push(() => runFileOperation(() => this.primeInputFile(file, operationContext))); } else if (file.content !== undefined) { - fileOps.push(this.writeFile(file)); + fileOps.push(() => runFileOperation(() => this.writeFile(file, operationContext))); + } + } + let nextOperation = 0; + const runPrimeWorker = async (): Promise => { + while (!controller.signal.aborted) { + const operationIndex = nextOperation++; + if (operationIndex >= fileOps.length) return; + await fileOps[operationIndex](); } + }; + const workerCount = Math.min(config.prime_concurrency, fileOps.length); + await Promise.allSettled( + Array.from({ length: workerCount }, () => runPrimeWorker()), + ); + if (firstFailure) { + if (this.session) { + /* A sibling may already have atomically replaced its destination. The + * workspace now matches neither the previous checkpoint nor the full + * request, so keeping this VM would silently preserve partial input + * state. Fail closed and give the control plane a stable recycle signal. */ + this.session.markDirty('input priming failed after file operations began'); + throw new SessionWorkspaceDirtyError(firstFailure.error); + } + throw firstFailure.error; + } + } + + /** + * Downloads an input file, or — in session mode when the same storage id is + * already present on disk from a prior call — skips the network fetch and + * hashes the local copy so modification detection still works. + */ + private async primeInputFile( + file: TFile, + context: PrimeOperationContext, + ): Promise { + throwIfAborted(context.signal); + if (this.session && file.id && (await this.reusePrimedInput(file, context))) { + /* Reuse has no response header to pass through downloadAndWriteFile, so + * its requested name becomes authoritative only after the on-disk copy + * has been verified. Reserve it before another concurrent ref can claim + * and overwrite that path. */ + this.reserveInputDestination(file, file.name); + return; + } + const name = await this.downloadAndWriteFile(file, 5, 500, context); + if (this.session && file.id) { + /* Record read-only so the next turn re-downloads it (primedInputId reports + * read-only primes as not-primed) — a reused on-disk copy could have been + * tampered via the writable parent dir. Keep the original upload hash as + * the reuse baseline so a later turn detects prior in-place mutations. */ + const primed = this.inputFileHashes.get(name); + this.session.markPrimed( + name, + this.inputIdentity(file), + primed?.readOnly === true, + primed?.hash, + ); + } + } + + private async reusePrimedInput( + file: TFile, + context?: PrimeOperationContext, + ): Promise { + const session = this.session; + if (!session || !file.id) return false; + throwIfAborted(context?.signal); + if (session.primedInputId(file.name) !== this.inputIdentity(file)) return false; + const submissionDir = context?.submissionDir ?? this.submissionDir; + const filePath = path.join(submissionDir, file.name); + /* lstat/O_NOFOLLOW protect only their final path component. A prior + * sandbox turn can instead replace an ancestor with a symlink, making an + * outside regular file look reusable and bypassing the guarded download + * path. Keep this outside the missing-file fallback catch: a symlinked + * ancestor is a hard session-integrity failure, not a cache miss. */ + await this.ensureDirNoFollow(path.dirname(filePath), submissionDir); + try { + const st = await fsp.lstat(filePath); + throwIfAborted(context?.signal); + if (!st.isFile()) return false; + /* Baseline against the ORIGINAL upload hash (recorded at prime time), not + * a re-hash of the on-disk copy: a prior turn may have mutated it in + * place, and re-hashing would bank the mutation as pristine and let the + * walker echo the original ref as unchanged. Falls back to hashing on + * disk if no original hash was retained. */ + const hash = this.session?.primedHash(file.name) ?? (await this.computeFileHash(filePath, true)); + this.inputFileHashes.set(file.name, { + originalId: file.id, + originalSessionId: file.storage_session_id, + hash, + path: filePath, + /* Carry the read-only bit into this run's input info: without it the + * output scan would treat a modified read-only input as a normal + * changed file and surface it, instead of dropping the modification + * per the read-only contract. */ + readOnly: session.isPrimedReadOnly(file.name) || undefined, + }); + return true; + } catch { + return false; } - await Promise.all(fileOps); } private fileEgressBaseUrl(): string { @@ -935,38 +1214,80 @@ export class Job { return true; } - async downloadAndWriteFile(file: TFile, maxRetries = 5, retryDelay = 500): Promise { - if (!file.id || !file.storage_session_id) return null; + async downloadAndWriteFile( + file: TFile, + maxRetries = 5, + retryDelay = 500, + context?: PrimeOperationContext, + ): Promise { + const operation: PrimeOperationContext = context ?? { + submissionDir: this.submissionDir, + identity: this.sandboxIdentity(), + }; + throwIfAborted(operation.signal); + if (!file.id || !file.storage_session_id) { + throw new ValidationError('By-reference inputs require id and storage_session_id'); + } - validateFilePath(file.name, this.submissionDir); + validateFilePath(file.name, operation.submissionDir); - const tempPath = path.join(this.submissionDir, `.tmp-${nanoid()}`); + const tempPath = path.join(operation.submissionDir, `.tmp-${nanoid()}`); let lastError: Error | null = null; for (let attempt = 1; attempt <= maxRetries; attempt++) { + let response: Response | undefined; try { - const response = await fetch(this.buildDownloadUrl(file), { - headers: this.fileEgressHeaders(), - }); + throwIfAborted(operation.signal); + response = await this.fetchInputObject(file, operation.signal); + if (operation.signal?.aborted) { + await response.body?.cancel().catch(() => {}); + throw abortReason(operation.signal); + } if (response.status === 404 && attempt < maxRetries) { + await response.body?.cancel().catch(() => {}); const delay = retryDelay * Math.pow(2, attempt - 1); this.log.info({ fileId: file.id, attempt, maxRetries, delay }, 'File not found, retrying'); - await sleep(delay); + await sleep(delay, operation.signal); continue; } - if (!response.ok) throw new Error(`HTTP error: ${response.status}`); + if (!response.ok) { + await response.body?.cancel().catch(() => {}); + throw new Error(`HTTP error: ${response.status}`); + } const originalName = resolveOriginalName(response, file); - validateFilePath(originalName, this.submissionDir); - const finalPath = path.join(this.submissionDir, originalName); + validateFilePath(originalName, operation.submissionDir); + this.reserveInputDestination(file, originalName); + const finalPath = path.join(operation.submissionDir, originalName); const finalParent = path.dirname(finalPath); - await fsp.mkdir(finalParent, { recursive: true }); - await this.secureAncestors(finalParent); + /* Persistent-session workspaces can hold a prior turn's symlink, so build + * ancestors no-follow; a fresh per-job workspace can use plain mkdir -p. */ + if (this.session) await this.ensureDirNoFollow(finalParent, operation.submissionDir); + else await fsp.mkdir(finalParent, { recursive: true }); + await this.secureAncestors(finalParent, operation.submissionDir, operation.identity); + /* Clear a symlink/dir a prior session turn may have squatted at the + * target so streamToDisk's rename lands a fresh regular file rather than + * following a link or failing on a directory. A regular file is LEFT in + * place and atomically replaced by that rename, so a failed/timed-out + * download can't delete prior session state before the replacement is + * safely written. */ + if (this.session) { + const existing = await fsp.lstat(finalPath).catch(() => null); + if (existing && !existing.isFile()) { + await fsp.rm(finalPath, { force: true, recursive: true }); + } + } - const hash = await this.streamToDisk(response, tempPath, finalPath); + const hash = await this.streamToDisk( + response, + tempPath, + finalPath, + operation.identity, + operation.signal, + ); const readOnly = response.headers.get('x-read-only')?.toLowerCase() === 'true'; this.inputFileHashes.set(originalName, { originalId: file.id, @@ -978,11 +1299,7 @@ export class Job { /* Defense-in-depth: keep read-only inputs root-owned + 0444 so the * sandbox UID can read them but cannot chmod them back to writable. */ if (readOnly) { - try { - await applyReadOnlyInputPermissions(finalPath); - } catch (err) { - this.log.warn({ file: originalName, err }, 'Failed to chmod read-only input'); - } + await applyReadOnlyInputPermissions(finalPath); } /* Keep the in-memory TFile in sync with the on-disk name so that @@ -995,6 +1312,13 @@ export class Job { this.log.info({ file: originalName, hash: hash.substring(0, 8) }, 'Downloaded file'); return originalName; } catch (error: unknown) { + if (response?.body && !response.bodyUsed) { + await response.body.cancel().catch(() => {}); + } + if (operation.signal?.aborted) { + try { await fsp.unlink(tempPath); } catch { /* may not exist */ } + throw abortReason(operation.signal); + } /* ValidationError is deterministic — a bad Content-Disposition * filename will fail identically on every retry. Abort fast * (cleanup + rethrow) instead of burning ~7.5s on exponential @@ -1007,14 +1331,55 @@ export class Job { if (attempt < maxRetries) { const delay = retryDelay * Math.pow(2, attempt - 1); this.log.warn({ fileId: file.id, attempt, maxRetries, delay, err: lastError }, 'Download failed, retrying'); - await sleep(delay); + await sleep(delay, operation.signal); } } } this.log.error({ fileId: file.id, maxRetries, err: lastError }, 'Failed to download file'); try { await fsp.unlink(tempPath); } catch { /* may not exist */ } - return null; + throw lastError ?? new Error(`Failed to download input ${file.id}`); + } + + /** + * Resolves an input object's bytes, preferring the runner-local cache the + * control plane pushes into on backends whose sandbox cannot reach the file + * server (see session-inputs.ts). A cache hit is presented as the very + * `Response` a fetch would have produced, so name resolution, read-only + * protection, hashing, ownership and priming all run identically for pushed + * and pulled inputs — there is exactly one workspace writer. + */ + private async fetchInputObject(file: TFile, signal?: AbortSignal): Promise { + throwIfAborted(signal); + const cached = await openCachedInput( + file.storage_session_id!, + file.id!, + file.input_cache_key, + ); + if (signal?.aborted) { + await cached?.handle.close().catch(() => {}); + throwIfAborted(signal); + } + if (cached) { + this.log.debug({ fileId: file.id }, 'Priming input from pushed cache'); + return cachedInputResponse(cached); + } + if (!this.fileEgressBaseUrl()) { + /* Push-model deployment with nothing pushed for this ref: fetching would + * hit an unreachable (or unset) file server and surface as a confusing + * transport error. Say what actually went wrong. */ + throw new Error( + `Input ${file.id} was not delivered to the sandbox and no file server is reachable`, + ); + } + return fetch(this.buildDownloadUrl(file), { + headers: this.fileEgressHeaders(), + signal, + }); + } + + private inputIdentity(file: TFile): string { + return file.input_cache_key ?? inputCacheKey(file.storage_session_id ?? '', file.id ?? ''); } /** @@ -1035,7 +1400,10 @@ export class Job { response: Response, tempPath: string, finalPath: string, + identity = this.sandboxIdentity(), + signal?: AbortSignal, ): Promise { + throwIfAborted(signal); const body = response.body; if (!body) throw new Error('Response body is null'); @@ -1045,25 +1413,55 @@ export class Job { }); const fileStream = fs.createWriteStream(tempPath, { mode: SANDBOX_FILE_MODE }); const reader = toNodeReadable(body); - await pipeline(reader, hashTransform, fileStream); - await fsp.rename(tempPath, finalPath); - await this.applySandboxFilePermissions(finalPath); + const abort = (): void => { + const error = abortReason(signal!); + reader.destroy(error); + hashTransform.destroy(error); + fileStream.destroy(error); + }; + signal?.addEventListener('abort', abort, { once: true }); + try { + throwIfAborted(signal); + await pipeline(reader, hashTransform, fileStream); + throwIfAborted(signal); + await fsp.rename(tempPath, finalPath); + await this.applySandboxFilePermissions(finalPath, false, identity); + } finally { + signal?.removeEventListener('abort', abort); + } return hashStream.digest('hex'); } - async writeFile(file: TFile): Promise { - validateFilePath(file.name, this.submissionDir); - const filePath = path.join(this.submissionDir, file.name); + async writeFile(file: TFile, context?: PrimeOperationContext): Promise { + const operation: PrimeOperationContext = context ?? { + submissionDir: this.submissionDir, + identity: this.sandboxIdentity(), + }; + throwIfAborted(operation.signal); + validateFilePath(file.name, operation.submissionDir); + const filePath = path.join(operation.submissionDir, file.name); const content = Buffer.from(file.content ?? '', (file.encoding as BufferEncoding) ?? 'utf8'); const parentDir = path.dirname(filePath); - await fsp.mkdir(parentDir, { recursive: true }); - await this.secureAncestors(parentDir); + if (this.session) await this.ensureDirNoFollow(parentDir, operation.submissionDir); + else await fsp.mkdir(parentDir, { recursive: true }); + await this.secureAncestors(parentDir, operation.submissionDir, operation.identity); + throwIfAborted(operation.signal); + /* In a persistent session workspace a prior turn could have left a symlink + * (or a directory) squatting this path; the default writeFile would follow + * the symlink and clobber its target as root. Remove whatever is there + * first (unlink never follows a link) so we always write a fresh regular + * file. A fresh per-job workspace has nothing here. */ + if (this.session) await fsp.rm(filePath, { force: true, recursive: true }); await fsp.writeFile(filePath, content); - await this.applySandboxFilePermissions(filePath); + await this.applySandboxFilePermissions(filePath, false, operation.identity); const hash = crypto.createHash('sha256').update(content).digest('hex'); this.inputFileHashes.set(file.name, { hash, path: filePath }); + /* Track inline source files as session inputs too (not just downloaded + * ones), so a later turn that switches entrypoint (e.g. main.py -> script.sh) + * doesn't re-surface the prior turn's source as a generated output. */ + if (this.session) this.session.markPrimed(file.name, file.id ?? '', false, hash); } async safeCall( @@ -1389,6 +1787,14 @@ export class Job { const id = nanoid(); this.sessionFiles.push({ id, name: keepPath, storage_session_id: this.outputSessionId }); this.generatedFiles.push({ id, name: keepPath, path: keepFullPath }); + /* A synthesized marker is a generated session output just like a regular + * file. Record its known empty-file digest so a successful upload commits + * it to the surfaced set; otherwise the persistent marker is re-uploaded + * on every later turn. */ + if (this.session) { + const emptyHash = crypto.createHash('sha256').update('').digest('hex'); + this.pendingSurfaced.set(id, { name: keepPath, signature: emptyHash }); + } return { collected: true, truncated: false }; } @@ -1482,7 +1888,8 @@ export class Job { /* Use lstat to stay consistent with classifyDirent's symlink filter — * following a symlink here would resurrect the exact escape vector * that the classification step already rejected. */ - size = (await fsp.lstat(fullPath)).size; + const st = await fsp.lstat(fullPath); + size = st.size; } catch (err) { this.log.debug({ path: relativePath, err }, 'walkDir: unable to stat file'); return { collected: false, truncated: false, stopLoop: false }; @@ -1491,20 +1898,51 @@ export class Job { return { collected: false, truncated: false, stopLoop: false }; } + /* Session mode output diffing + input-modification detection. Hash by + * CONTENT (not size+mtime): a program can rewrite a surfaced output with + * different bytes while preserving size+mtime (os.utime / touch -r), which a + * stat-only signature would wrongly suppress. Compute once per session/input + * file and reuse for the suppression check, wasModified, and the surfaced + * mark; non-session jobs still only hash their inputs. */ const inputFileInfo = this.inputFileHashes.get(relativePath); const existingFile = inputByName.get(relativePath); - let wasModified = false; - - if (inputFileInfo) { + let contentHash: string | undefined; + if (inputFileInfo != null || this.session != null) { try { - const currentHash = await this.computeFileHash(fullPath, true); - wasModified = currentHash !== inputFileInfo.hash; - if (wasModified) this.log.info({ file: relativePath }, 'Input file was modified'); + contentHash = await this.computeFileHash(fullPath, true); } catch (err) { this.log.debug({ path: relativePath, err }, 'walkDir: failed to hash file'); } } + /* Skip a pure output already surfaced with identical content. */ + if (this.session && inputFileInfo == null && contentHash != null + && this.session.isSurfaced(relativePath, contentHash)) { + return { collected: false, truncated: false, stopLoop: false }; + } + + /* Skip a file primed as an input on an earlier turn that this turn didn't + * re-send: it has no inputFileInfo and was never surfaced as an output, so + * it would otherwise be echoed as a brand-new generated file on unrelated + * executions. Only suppress it while UNCHANGED vs the primed baseline: if + * the sandbox modified it since, surface the new bytes (also lets the + * upload-retry path recover a modified input whose earlier upload was + * pruned). Read-only inputs never surface modifications (dropped by + * contract), so always suppress those. */ + if (this.session && inputFileInfo == null && this.session.isPrimedInput(relativePath)) { + const primedHash = this.session.primedHash(relativePath); + const unchanged = contentHash != null && primedHash != null && contentHash === primedHash; + if (unchanged || this.session.isPrimedReadOnly(relativePath)) { + return { collected: false, truncated: false, stopLoop: false }; + } + } + + let wasModified = false; + if (inputFileInfo && contentHash != null) { + wasModified = contentHash !== inputFileInfo.hash; + if (wasModified) this.log.info({ file: relativePath }, 'Input file was modified'); + } + const echoed = this.tryEchoUnchangedInput({ wasModified, inputFileInfo, @@ -1528,6 +1966,12 @@ export class Job { } this.sessionFiles.push(fileData); this.generatedFiles.push({ id: newId, name: relativePath, path: fullPath }); + /* Defer the surfaced-mark until the upload succeeds (flushed in + * uploadGeneratedFiles) — marking here would suppress the file forever if + * its upload later fails and the route prunes it from the response. */ + if (this.session && contentHash != null) { + this.pendingSurfaced.set(newId, { name: relativePath, signature: contentHash }); + } return { collected: true, truncated: false, stopLoop: false }; } @@ -1694,6 +2138,15 @@ export class Job { if (id) uploaded.add(id); } + /* Commit the surfaced-mark only for files that actually uploaded, so a + * dropped upload leaves the file eligible to surface again next turn. */ + if (this.session) { + for (const id of uploaded) { + const pending = this.pendingSurfaced.get(id); + if (pending) this.session.markSurfaced(pending.name, pending.signature); + } + } + if (uploaded.size < this.generatedFiles.length) { this.log.warn( { uploaded: uploaded.size, total: this.generatedFiles.length }, @@ -1745,11 +2198,12 @@ export class Job { } const url = `${this.fileEgressBaseUrl()}/sessions/${encodeURIComponent(this.outputSessionId)}/objects/${encodeURIComponent(file.id)}`; - const controller = new AbortController(); - const timeout = setTimeout(() => controller.abort(), 30000); - let headers: Record; + let timeout: ReturnType | undefined; + let uploadHandle: Awaited> | undefined; + let stream: fs.ReadStream | undefined; + let response: Response | undefined; try { - headers = this.fileEgressHeaders({ + const headers = this.fileEgressHeaders({ /* file-server URL-decodes this header to recover the canonical * filename, so paths with `/` survive transport without colliding * with the `___` separators or RFC 5987 quoting rules used @@ -1763,19 +2217,14 @@ export class Job { 'Content-Type': mimeTypeFor(file.name), 'Content-Length': String(size), }); - } catch (error) { - clearTimeout(timeout); - this.log.error({ file: file.name, err: error }, 'Error preparing upload'); - return null; - } - - const stream = fs.createReadStream(file.path, { flags: fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW }); - stream.on('error', (error) => { - this.log.warn({ file: file.name, err: error }, 'Upload file stream error'); - }); - - let response: Response | undefined; - try { + /* See computeFileHash: numeric O_NOFOLLOW is only typed via fsp.open. */ + uploadHandle = await fsp.open(file.path, fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW); + stream = uploadHandle.createReadStream(); + stream.on('error', (error) => { + this.log.warn({ file: file.name, err: error }, 'Upload file stream error'); + }); + const controller = new AbortController(); + timeout = setTimeout(() => controller.abort(), 30000); response = await fetch(url, { method: 'PUT', headers, @@ -1798,8 +2247,9 @@ export class Job { this.log.error({ file: file.name, err: error }, 'Error uploading file'); return null; } finally { - clearTimeout(timeout); - stream.destroy(); + if (timeout) clearTimeout(timeout); + stream?.destroy(); + await uploadHandle?.close().catch(() => {}); /* Drain or cancel the response body. Undici keeps the socket * reserved until the body is consumed; under concurrent uploads, * leaving bodies unread exhausts the connection pool and stalls @@ -1818,6 +2268,17 @@ export class Job { if (!this.isSynthetic) { this.log.info('Cleaning up'); } + + /* Session mode: the workspace and pinned UID belong to the long-lived + * session, not this job. Keep both so the next call sees prior files; + * teardown happens on the /terminate hook (or explicit session reset). */ + if (this.session) { + this.workspaceLease = undefined; + this.submissionDir = ''; + this.jobIdentity = undefined; + return; + } + let workspaceRemoved = true; const workspaceLease = this.workspaceLease; const jobIdentity = this.jobIdentity; @@ -1855,6 +2316,28 @@ export class Job { } } -function sleep(ms: number): Promise { - return new Promise(resolve => setTimeout(resolve, ms)); +function abortReason(signal: AbortSignal): Error { + return signal.reason instanceof Error + ? signal.reason + : new Error('Input priming aborted'); +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) throw abortReason(signal); +} + +function sleep(ms: number, signal?: AbortSignal): Promise { + if (!signal) return new Promise(resolve => setTimeout(resolve, ms)); + throwIfAborted(signal); + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + signal.removeEventListener('abort', abort); + resolve(); + }, ms); + const abort = (): void => { + clearTimeout(timer); + reject(abortReason(signal)); + }; + signal.addEventListener('abort', abort, { once: true }); + }); } diff --git a/api/src/nsjail.test.ts b/api/src/nsjail.test.ts index f19d53a..ee6fbbc 100644 --- a/api/src/nsjail.test.ts +++ b/api/src/nsjail.test.ts @@ -187,6 +187,119 @@ exit 0 await fsp.rm(tmp, { recursive: true, force: true }); } }); + + test('stderr overflow truncates with an in-band marker instead of killing the job', async () => { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'nsjail-stderr-overflow-')); + const fakeNsJail = path.join(tmp, 'fake-nsjail.sh'); + const cfg = path.join(tmp, 'sandbox.cfg'); + const submissionDir = path.join(tmp, 'submission'); + await fsp.mkdir(submissionDir); + await fsp.writeFile(cfg, ''); + await fsp.writeFile( + fakeNsJail, + `#!/bin/sh +log_path="" +while [ "$#" -gt 0 ]; do + if [ "$1" = "--log" ]; then + shift + log_path="$1" + fi + shift || break +done +if [ -n "$log_path" ]; then + printf '[I][test] Executing "/bin/bash" for noisy child\\n' > "$log_path" +fi +head -c 3000 /dev/zero | tr '\\0' 'e' 1>&2 +printf 'survived the noise\\n' +exit 0 +`, + { mode: 0o755 }, + ); + + const originalNsJailPath = config.nsjail_path; + const originalNsJailConfig = config.nsjail_config; + config.nsjail_path = fakeNsJail; + config.nsjail_config = cfg; + try { + const result = await execute({ + command: ['/bin/bash', '/pkgs/bash/5.2.0/run', 'main.sh'], + envVars: {}, + submissionDir, + pkgdir: '/pkgs/bash/5.2.0', + timeout: 5000, + memoryLimit: -1, + outputMaxSize: 1024, + identity: { slot: 0, uid: 65534, gid: 65534, perJobUid: false }, + }); + + expect(result.code).toBe(0); + expect(result.signal).toBeNull(); + expect(result.stdout).toBe('survived the noise\n'); + expect(result.message).toBe('stderr truncated at 1024 bytes'); + expect(result.stderr.endsWith('\n[sandbox] stderr truncated at 1024 bytes\n')).toBe(true); + expect(result.stderr.length).toBeLessThanOrEqual( + 1024 + '\n[sandbox] stderr truncated at 1024 bytes\n'.length, + ); + } finally { + config.nsjail_path = originalNsJailPath; + config.nsjail_config = originalNsJailConfig; + await fsp.rm(tmp, { recursive: true, force: true }); + } + }); + + test('kill reasons surface as an in-band [sandbox] marker on stderr', async () => { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'nsjail-kill-marker-')); + const fakeNsJail = path.join(tmp, 'fake-nsjail.sh'); + const cfg = path.join(tmp, 'sandbox.cfg'); + const submissionDir = path.join(tmp, 'submission'); + await fsp.mkdir(submissionDir); + await fsp.writeFile(cfg, ''); + await fsp.writeFile( + fakeNsJail, + `#!/bin/sh +log_path="" +while [ "$#" -gt 0 ]; do + if [ "$1" = "--log" ]; then + shift + log_path="$1" + fi + shift || break +done +if [ -n "$log_path" ]; then + printf '[I][test] Executing "/bin/bash" for slow child\\n' > "$log_path" + printf '[W][test] run time >= time limit\\n' >> "$log_path" +fi +exit 137 +`, + { mode: 0o755 }, + ); + + const originalNsJailPath = config.nsjail_path; + const originalNsJailConfig = config.nsjail_config; + config.nsjail_path = fakeNsJail; + config.nsjail_config = cfg; + try { + const result = await execute({ + command: ['/bin/bash', '/pkgs/bash/5.2.0/run', 'main.sh'], + envVars: {}, + submissionDir, + pkgdir: '/pkgs/bash/5.2.0', + timeout: 5000, + memoryLimit: -1, + outputMaxSize: 1024, + identity: { slot: 0, uid: 65534, gid: 65534, perJobUid: false }, + }); + + expect(result.status).toBe('TO'); + expect(result.message).toBe('Time limit exceeded'); + expect(result.stderr.endsWith('\n[sandbox] Time limit exceeded\n')).toBe(true); + expect(result.output.endsWith('\n[sandbox] Time limit exceeded\n')).toBe(true); + } finally { + config.nsjail_path = originalNsJailPath; + config.nsjail_config = originalNsJailConfig; + await fsp.rm(tmp, { recursive: true, force: true }); + } + }); }); describe('renderJobConfigOverlay', () => { diff --git a/api/src/nsjail.ts b/api/src/nsjail.ts index c07a384..87a2b79 100644 --- a/api/src/nsjail.ts +++ b/api/src/nsjail.ts @@ -272,6 +272,7 @@ export async function execute(opts: ExecuteOptions, setupGate: NsJailSetupGate = let killed = false; let killMessage: string | null = null; let killStatus: string | null = null; + let stderrTruncated = false; function drainStream( child: ChildProcessWithoutNullStreams, @@ -302,18 +303,23 @@ export async function execute(opts: ExecuteOptions, setupGate: NsJailSetupGate = stdout += chunk; output += chunk; } else { + /* stderr overflow must not kill the job: noisy-but-healthy + * processes (e.g. libraries that dump stack traces per operation) + * would lose their entire computation over diagnostics. Truncate + * at the cap and keep draining so the pipe never back-pressures; + * a marker is appended after exit so consumers see what happened. + * stdout overflow still kills below — stdout is the product, and + * a runaway producer should stop paying for compute. */ + if (stderrTruncated) { + return; + } if (stderr.length + chunk.length > outputMaxSize) { const remaining = outputMaxSize - stderr.length; if (remaining > 0) { stderr += chunk.slice(0, remaining); output += chunk.slice(0, remaining); } - if (!killed) { - killMessage = 'stderr length exceeded'; - killStatus = 'EL'; - killed = true; - child.kill('SIGKILL'); - } + stderrTruncated = true; return; } stderr += chunk; @@ -583,13 +589,26 @@ export async function execute(opts: ExecuteOptions, setupGate: NsJailSetupGate = signal = SIGNALS[code - 128] ?? null; } - const finalMessage = killMessage ?? logMessage; + let finalMessage = killMessage ?? logMessage; const finalStatus = killStatus ?? logStatus; if (finalStatus && ['TO', 'OL', 'EL'].includes(finalStatus)) { signal = 'SIGKILL'; } + if (stderrTruncated) { + finalMessage = finalMessage ?? `stderr truncated at ${outputMaxSize} bytes`; + } + + /* Surface the reason in the streams themselves: most clients (and the + * models consuming tool output) only ever see stdout/stderr, so a kill + * or truncation without an in-band marker reads as a silent hang. */ + if (finalMessage) { + const marker = `\n[sandbox] ${finalMessage}\n`; + stderr += marker; + output += marker; + } + return { stdout, stderr, diff --git a/api/src/session-checkpoint.test.ts b/api/src/session-checkpoint.test.ts new file mode 100644 index 0000000..d0eecef --- /dev/null +++ b/api/src/session-checkpoint.test.ts @@ -0,0 +1,817 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { spawnSync } from 'child_process'; +import { PassThrough, Readable } from 'stream'; +import { gzipSync, gunzipSync } from 'zlib'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { config } from './config'; +import type { SandboxJobIdentity } from './workspace-isolation'; +import type { SessionWorkspace } from './session-workspace'; +import { SANDBOX_WORKSPACE_ROOT, SESSION_WORKSPACE_ID, fallbackSandboxIdentity } from './workspace-isolation'; +import { restoreSessionCheckpoint, streamSessionCheckpoint } from './session-checkpoint'; +import { + SESSION_META_FILE, + SESSION_META_MARKER, + bindSessionWorkspace, + resetSessionWorkspaceStateForTests, + unbindSessionWorkspace, +} from './session-workspace'; + +const savedEnabled = config.session_workspace_enabled; +const savedPerJob = config.per_job_uids; +const savedCheckpointMaxBytes = config.checkpoint_max_bytes; +const CHECKPOINT_CONTROL_FILE = '.codeapi-checkpoint-control.v2.json'; +const CHECKPOINT_CONTROL_MAX_BYTES = 16 * 1024 * 1024; + +afterEach(async () => { + config.session_workspace_enabled = savedEnabled; + config.per_job_uids = savedPerJob; + config.checkpoint_max_bytes = savedCheckpointMaxBytes; + await unbindSessionWorkspace().catch(() => {}); + resetSessionWorkspaceStateForTests(); + await fsp + .rm(path.join(SANDBOX_WORKSPACE_ROOT, SESSION_WORKSPACE_ID), { recursive: true, force: true }) + .catch(() => {}); +}); + +/** CI and local dev run bun as a non-root user, where the default per-job-UID + * configuration requires root for workspace chowns. Switch the session to the + * shared fallback identity (perJobUid=false) and flip the config flag the + * workspace-root preparation consults, so skipped chowns degrade to + * compatibility modes — the same degradation the runner itself applies when + * running unprivileged outside hardened mode. */ +function seedNonRootIdentity(session: SessionWorkspace): void { + config.per_job_uids = false; + (session as unknown as { identity?: SandboxJobIdentity }).identity = fallbackSandboxIdentity(); +} + +/** Minimal Express response double capturing status + json body. */ +function fakeRes(): { status: number; body: unknown; setHeader: () => void; destroy: () => void } & { + status(code: number): { json(body: unknown): void }; +} { + const res = { + statusCode: 0, + body: undefined as unknown, + setHeader: () => {}, + destroy: () => {}, + status(code: number) { + res.statusCode = code; + return { + json(body: unknown) { res.body = body; }, + }; + }, + }; + return res as never; +} + +function fakeStreamRes(): { statusCode: number; body: unknown; headersSent: boolean } & { + status(code: number): { json(body: unknown): void }; +} { + const res = { + statusCode: 0, + body: undefined as unknown, + headersSent: false, + status(code: number) { + res.statusCode = code; + return { + json(body: unknown) { res.body = body; res.headersSent = true; }, + }; + }, + }; + return res as never; +} + +type CheckpointCaptureResponse = PassThrough & { + statusCode: number; + body: unknown; + headersSent: boolean; + status(code: number): CheckpointCaptureResponse; + json(body: unknown): CheckpointCaptureResponse; + setHeader(): CheckpointCaptureResponse; +}; + +function checkpointCaptureRes(): { + res: CheckpointCaptureResponse; + archive(): Buffer; +} { + const chunks: Buffer[] = []; + const res = new PassThrough() as CheckpointCaptureResponse; + res.statusCode = 0; + res.body = undefined; + res.headersSent = false; + res.status = (code: number) => { + res.statusCode = code; + return res; + }; + res.json = (body: unknown) => { + res.body = body; + res.headersSent = true; + res.end(); + return res; + }; + res.setHeader = () => res; + res.on('data', (chunk: Buffer) => { + res.headersSent = true; + chunks.push(Buffer.from(chunk)); + }); + return { res, archive: () => Buffer.concat(chunks) }; +} + +function readArchiveMember(archive: Buffer, member: string): Buffer { + const extracted = spawnSync( + 'tar', + ['-xOzf', '-', member], + { input: archive, maxBuffer: 64 * 1024 * 1024 }, + ); + if (extracted.status !== 0) { + throw new Error(`fixture tar extraction exited ${extracted.status}`); + } + return extracted.stdout; +} + +/** Builds a real tar.gz whose members live under a leading `session/` dir, + * matching the archive shape the checkpoint create side produces. */ +async function makeArchive( + files: Record, + topLevelFiles: Record = {}, +): Promise { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'sess-ckpt-')); + const stage = path.join(tmp, 'session'); + for (const [name, content] of Object.entries(files)) { + const target = path.join(stage, name); + await fsp.mkdir(path.dirname(target), { recursive: true }); + await fsp.writeFile(target, content); + } + for (const [name, content] of Object.entries(topLevelFiles)) { + await fsp.writeFile(path.join(tmp, name), content); + } + const tar = spawnSync( + 'tar', + ['-czf', '-', '-C', tmp, 'session', ...Object.keys(topLevelFiles)], + { maxBuffer: 64 * 1024 * 1024 }, + ); + await fsp.rm(tmp, { recursive: true, force: true }); + if (tar.status !== 0) throw new Error(`fixture tar exited ${tar.status}`); + return tar.stdout; +} + +describe('session checkpoint gating', () => { + test('checkpoint is 409 when no session is bound', async () => { + const res = fakeRes(); + await streamSessionCheckpoint(res as never); + expect((res as unknown as { statusCode: number }).statusCode).toBe(409); + }); + + test('restore is 409 when no session is bound', async () => { + const res = fakeRes(); + await restoreSessionCheckpoint({} as never, res as never); + expect((res as unknown as { statusCode: number }).statusCode).toBe(409); + }); + + test('checkpoint is 409 while the bound workspace is dirty', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_dirty' }); + seedNonRootIdentity(session!); + session!.markDirty('partial input delivery'); + + const res = fakeRes(); + await streamSessionCheckpoint(res as never); + + expect((res as unknown as { statusCode: number }).statusCode).toBe(409); + expect(res.body).toEqual({ + error: 'session_workspace_dirty', + message: 'Session workspace must be restored before checkpointing', + }); + }); +}); + +describe('streamSessionCheckpoint', () => { + test('does not publish a clean response EOF when tar fails after emitting bytes', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_late_tar_failure' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'checkpoint.txt'), 'checkpoint bytes'); + const fakeTarDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'late-tar-failure-')); + const fakeTar = path.join(fakeTarDir, 'tar'); + await fsp.writeFile( + fakeTar, + '#!/bin/sh\nprintf "partial tar bytes"\nexit 23\n', + { mode: 0o700 }, + ); + + try { + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never, { tarCommand: fakeTar }); + + expect(capture.archive().length).toBeGreaterThan(0); + expect(capture.res.headersSent).toBe(true); + expect(capture.res.writableEnded).toBe(false); + expect(capture.res.destroyed).toBe(true); + } finally { + await fsp.rm(fakeTarDir, { recursive: true, force: true }); + } + }); + + test('rejects creation when the expanded tar exceeds the restore cap', async () => { + config.session_workspace_enabled = true; + config.checkpoint_max_bytes = 32 * 1024; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_expanded_cap' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + const highlyCompressible = 'x'.repeat(64 * 1024); + await fsp.writeFile(path.join(dir, 'highly-compressible.txt'), highlyCompressible); + + /* This is the precise failure mode: the gzip is comfortably accepted by + * the compressed-byte cap even though restore must reject its expanded tar + * stream. Create must reject it before publishing a recovery point. */ + const compressedFixture = await makeArchive({ + 'highly-compressible.txt': highlyCompressible, + }); + expect(compressedFixture.length).toBeLessThan(config.checkpoint_max_bytes); + + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never); + + expect(capture.res.headersSent).toBe(true); + expect(capture.res.destroyed).toBe(true); + }); + + test('a checkpoint accepted by both create-side caps round-trips under the same cap', async () => { + config.session_workspace_enabled = true; + config.checkpoint_max_bytes = 64 * 1024; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_symmetric_cap' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + const checkpointBytes = 'x'.repeat(24 * 1024); + await fsp.writeFile(path.join(dir, 'roundtrip.txt'), checkpointBytes); + + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never); + + expect(capture.res.writableEnded).toBe(true); + expect(capture.res.writableFinished).toBe(true); + const archive = capture.archive(); + expect(archive.length).toBeLessThanOrEqual(config.checkpoint_max_bytes); + expect(gunzipSync(archive).length).toBeLessThanOrEqual(config.checkpoint_max_bytes); + + await fsp.writeFile(path.join(dir, 'roundtrip.txt'), 'stale'); + const restoreRes = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(archive) as never, + restoreRes as never, + ); + + expect(restoreRes.statusCode).toBe(200); + expect(await fsp.readFile(path.join(dir, 'roundtrip.txt'), 'utf8')).toBe(checkpointBytes); + }); + + test('retains accumulated surfaced state while the full control record fits', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_surfaced_cap' }); + seedNonRootIdentity(session!); + await session!.ownership(); + session!.markPrimed('input.csv', 'stable-cache-key', false, 'original-hash'); + for (let i = 0; i < 4_000; i++) { + const suffix = `${i}-${'s'.repeat(80)}`; + session!.markSurfaced(`generated/${suffix}.txt`, `signature-${suffix}`); + } + + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never); + + expect(capture.res.statusCode).toBe(200); + const control = JSON.parse( + readArchiveMember(capture.archive(), CHECKPOINT_CONTROL_FILE).toString(), + ); + expect(control.marker).toBe(SESSION_META_MARKER); + expect(control.primed).toEqual([[ + 'input.csv', + { id: 'stable-cache-key', readOnly: false, hash: 'original-hash' }, + ]]); + expect(control.surfaced).toHaveLength(4_000); + }); + + test('round-trips control state larger than the old 256 KiB ceiling', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_control_cap' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'current.txt'), 'current workspace bytes'); + + /* Both maps exceed the old control budget. The bounded parser can safely + * carry this recovery state, so neither map should be degraded. */ + for (let i = 0; i < 4_000; i++) { + const suffix = `${i}-${'s'.repeat(80)}`; + session!.markSurfaced(`generated/${suffix}.txt`, `signature-${suffix}`); + session!.markPrimed( + `inputs/${suffix}.txt`, + `cache-key-${suffix}`, + false, + `hash-${suffix}`, + ); + } + expect( + Buffer.byteLength(JSON.stringify({ + marker: SESSION_META_MARKER, + ...session!.snapshotMeta(), + })), + ).toBeGreaterThan(256 * 1024); + expect( + Buffer.byteLength(JSON.stringify({ + marker: SESSION_META_MARKER, + primed: session!.snapshotMeta().primed, + surfaced: [], + })), + ).toBeGreaterThan(256 * 1024); + + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never); + + expect(capture.res.statusCode).toBe(200); + const archive = capture.archive(); + expect(readArchiveMember(archive, 'session/current.txt').toString()) + .toBe('current workspace bytes'); + const control = readArchiveMember(archive, CHECKPOINT_CONTROL_FILE); + const parsed = JSON.parse(control.toString()); + expect(parsed.marker).toBe(SESSION_META_MARKER); + expect(parsed.surfaced).toHaveLength(4_000); + expect(parsed.primed).toHaveLength(4_000); + + /* Prove the recovery point rebuilds the primed identity used by the next + * execute to preserve an in-place-modified input instead of downloading + * the original over it. */ + session!.loadMeta({ primed: [], surfaced: [] }); + const restoreRes = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(archive) as never, + restoreRes as never, + ); + expect(restoreRes.statusCode).toBe(200); + expect(session!.primedInputId(`inputs/0-${'s'.repeat(80)}.txt`)) + .toBe(`cache-key-0-${'s'.repeat(80)}`); + }); + + test('drops surfaced state only beyond the hard cap and preserves primed identity', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_hard_cap' }); + seedNonRootIdentity(session!); + await session!.ownership(); + session!.markPrimed('input.csv', 'stable-cache-key', false, 'original-hash'); + session!.markSurfaced('generated/huge.txt', 's'.repeat(CHECKPOINT_CONTROL_MAX_BYTES)); + + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never); + + expect(capture.res.statusCode).toBe(200); + const control = JSON.parse( + readArchiveMember(capture.archive(), CHECKPOINT_CONTROL_FILE).toString(), + ); + expect(control).toEqual({ + marker: SESSION_META_MARKER, + primed: [[ + 'input.csv', + { id: 'stable-cache-key', readOnly: false, hash: 'original-hash' }, + ]], + surfaced: [], + }); + }); +}); + +describe('restoreSessionCheckpoint', () => { + test('round-trips control metadata without touching a user file at the legacy path', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_checkpoint_roundtrip' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + const userFile = JSON.stringify({ + marker: SESSION_META_MARKER, + primed: [['user-trap.csv', { + id: 'user-controlled-id', + readOnly: false, + }]], + surfaced: [], + }); + await fsp.writeFile(path.join(dir, SESSION_META_FILE), userFile); + await fsp.writeFile(path.join(dir, 'session-meta.json'), 'rollback-sensitive-user-data'); + await fsp.writeFile(path.join(dir, 'input.csv'), 'sandbox-modified'); + session!.markPrimed('input.csv', 'stable-cache-key', false, 'original-hash'); + session!.markSurfaced('output.csv', 'output-hash'); + + const capture = checkpointCaptureRes(); + await streamSessionCheckpoint(capture.res as never); + + expect(capture.res.statusCode).toBe(200); + expect(await fsp.readFile(path.join(dir, SESSION_META_FILE), 'utf8')).toBe(userFile); + const archive = capture.archive(); + + /* A pre-control-metadata image restores with --strip-components=1. The + * runner control member must be a single top-level component so that old + * tar skips it instead of extracting/overwriting `session-meta.json`. */ + const legacyRestore = await fsp.mkdtemp(path.join(os.tmpdir(), 'legacy-restore-')); + try { + const legacyTar = spawnSync( + 'tar', + ['-xzf', '-', '--strip-components=1', '-C', legacyRestore], + { input: archive, maxBuffer: 64 * 1024 * 1024 }, + ); + expect(legacyTar.status).toBe(0); + expect(await fsp.readFile(path.join(legacyRestore, 'session-meta.json'), 'utf8')) + .toBe('rollback-sensitive-user-data'); + } finally { + await fsp.rm(legacyRestore, { recursive: true, force: true }); + } + + session!.loadMeta({ primed: [], surfaced: [] }); + const restoreRes = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(archive) as never, + restoreRes as never, + ); + + expect(restoreRes.statusCode).toBe(200); + expect(await fsp.readFile(path.join(dir, SESSION_META_FILE), 'utf8')).toBe(userFile); + expect(session!.primedInputId('input.csv')).toBe('stable-cache-key'); + expect(session!.primedInputId('user-trap.csv')).toBeUndefined(); + expect(session!.isSurfaced('output.csv', 'output-hash')).toBe(true); + }); + + test('replaces the workspace with the archive contents', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_1' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + /* Restore is a full replace, unlike input delivery: state predating the + * checkpoint must not survive it. */ + await fsp.writeFile(path.join(dir, 'stale.txt'), 'from-a-previous-life'); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive({ 'restored.csv': 'a,b\n1,2\n' })) as never, + res as never, + ); + + expect(res.statusCode).toBe(200); + expect(await fsp.readFile(path.join(dir, 'restored.csv'), 'utf8')).toBe('a,b\n1,2\n'); + expect(await fsp.lstat(path.join(dir, 'stale.txt')).catch(() => null)).toBeNull(); + }); + + test('rolls the live workspace back when staged installation fails', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_install_failure' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'original live bytes'); + session!.markPrimed('live.txt', 'original-cache-key'); + let renameCalls = 0; + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive({ 'restored.txt': 'new checkpoint bytes' })) as never, + res as never, + { + rename: async (source, destination) => { + renameCalls += 1; + if (renameCalls === 2) throw new Error('injected staged install failure'); + await fsp.rename(source, destination); + }, + }, + ); + + expect(renameCalls).toBe(3); + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')).toBe('original live bytes'); + expect(await fsp.lstat(path.join(dir, 'restored.txt')).catch(() => null)).toBeNull(); + expect(session!.primedInputId('live.txt')).toBe('original-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('retains the live backup and marks dirty when installation rollback also fails', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_rollback_failure' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'recoverable live bytes'); + const stagesBefore = new Set( + (await fsp.readdir(SANDBOX_WORKSPACE_ROOT)) + .filter(name => name.startsWith('.session-restore-')), + ); + let renameCalls = 0; + let retainedStage: string | undefined; + + try { + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive({ 'restored.txt': 'new checkpoint bytes' })) as never, + res as never, + { + rename: async (source, destination) => { + renameCalls += 1; + if (renameCalls >= 2) throw new Error(`injected rename failure ${renameCalls}`); + await fsp.rename(source, destination); + }, + }, + ); + + retainedStage = (await fsp.readdir(SANDBOX_WORKSPACE_ROOT)) + .filter(name => name.startsWith('.session-restore-') && !stagesBefore.has(name)) + .map(name => path.join(SANDBOX_WORKSPACE_ROOT, name))[0]; + expect(renameCalls).toBe(3); + expect(res.statusCode).toBe(500); + expect(session!.dirtyReason).toBe('checkpoint restore rollback failed'); + expect(await fsp.readFile( + path.join(retainedStage!, '.live-workspace-backup', 'live.txt'), + 'utf8', + )).toBe('recoverable live bytes'); + } finally { + if (retainedStage) { + await fsp.rm(retainedStage, { recursive: true, force: true }); + } + } + }); + + test('keeps the committed restore when response delivery fails', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_response_failure' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'old live bytes'); + session!.markPrimed('live.txt', 'old-cache-key'); + const responseFailure = new Error('injected response failure'); + const res = { + statusCode: 0, + headersSent: false, + status(code: number) { + res.statusCode = code; + return { + json() { + res.headersSent = true; + throw responseFailure; + }, + }; + }, + }; + const restoredMeta = JSON.stringify({ + marker: SESSION_META_MARKER, + primed: [['restored.txt', { + id: 'restored-cache-key', + readOnly: false, + hash: 'restored-hash', + }]], + surfaced: [], + }); + + await expect(restoreSessionCheckpoint( + Readable.from(await makeArchive( + { 'restored.txt': 'new checkpoint bytes' }, + { [CHECKPOINT_CONTROL_FILE]: restoredMeta }, + )) as never, + res as never, + )).rejects.toThrow('injected response failure'); + + expect(res.statusCode).toBe(200); + expect(await fsp.readFile(path.join(dir, 'restored.txt'), 'utf8')) + .toBe('new checkpoint bytes'); + expect(await fsp.lstat(path.join(dir, 'live.txt')).catch(() => null)).toBeNull(); + expect(session!.primedInputId('live.txt')).toBeUndefined(); + expect(session!.primedInputId('restored.txt')).toBe('restored-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('consumes tar padding through expanded EOF before closing tar stdin', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_trailing_padding' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + const archive = await makeArchive({ 'restored.csv': 'a,b\n1,2\n' }); + const paddedArchive = gzipSync(Buffer.concat([ + gunzipSync(archive), + Buffer.alloc(128 * 1024), + ]), { level: 0 }); + + /* A decompressor may still have valid tar padding to emit after tar sees + * the archive's zero end markers. An uncompressed gzip stores a complete + * 64 KiB block in the first chunk below, including the tar end markers, + * while the delayed remainder keeps the expanded stream open. */ + async function* delayedArchive(): AsyncGenerator { + yield paddedArchive.subarray(0, 70 * 1024); + await new Promise(resolve => setTimeout(resolve, 100)); + yield paddedArchive.subarray(70 * 1024); + } + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(delayedArchive()) as never, + res as never, + ); + + expect(res.statusCode).toBe(200); + expect(await fsp.readFile(path.join(dir, 'restored.csv'), 'utf8')).toBe('a,b\n1,2\n'); + }); + + test('rejects a checkpoint upload that exceeds the runner-local compressed-byte cap', async () => { + config.session_workspace_enabled = true; + config.checkpoint_max_bytes = 64; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_oversized_upload' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'preserve compressed-cap state'); + session!.markPrimed('live.txt', 'live-cache-key'); + const archive = await makeArchive({ 'restored.csv': 'new bytes' }); + expect(archive.length).toBeGreaterThan(config.checkpoint_max_bytes); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint(Readable.from(archive) as never, res as never); + + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')) + .toBe('preserve compressed-cap state'); + expect(session!.primedInputId('live.txt')).toBe('live-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('rejects a checkpoint whose expanded tar stream exceeds the runner-local cap', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_expansion_bomb' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'preserve expanded-cap state'); + session!.markPrimed('live.txt', 'live-cache-key'); + const archive = await makeArchive({ 'highly-compressible.txt': 'x'.repeat(32 * 1024) }); + config.checkpoint_max_bytes = archive.length + 1024; + expect(archive.length).toBeLessThan(config.checkpoint_max_bytes); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint(Readable.from(archive) as never, res as never); + + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')) + .toBe('preserve expanded-cap state'); + expect(session!.primedInputId('live.txt')).toBe('live-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('rejects malformed metadata in a present new-format control member', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_bad_control' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'preserve malformed-control state'); + session!.markPrimed('stale.csv', 'stale-cache-key'); + session!.markDirty('pre-existing dirty state'); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive( + { 'input.csv': 'sandbox-modified' }, + { [CHECKPOINT_CONTROL_FILE]: '{"marker":' }, + )) as never, + res as never, + ); + + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')) + .toBe('preserve malformed-control state'); + expect(session!.primedInputId('stale.csv')).toBe('stale-cache-key'); + expect(session!.dirtyReason).toBe('pre-existing dirty state'); + }); + + test('rejects new-format control metadata larger than 16 MiB before parsing it', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_large_control' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'preserve oversized-control state'); + session!.markPrimed('live.txt', 'live-cache-key'); + const validButOversized = `${' '.repeat(CHECKPOINT_CONTROL_MAX_BYTES)}${JSON.stringify({ + marker: SESSION_META_MARKER, + primed: [], + surfaced: [], + })}`; + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive( + { 'input.csv': 'sandbox-modified' }, + { [CHECKPOINT_CONTROL_FILE]: validButOversized }, + )) as never, + res as never, + ); + + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')) + .toBe('preserve oversized-control state'); + expect(session!.primedInputId('live.txt')).toBe('live-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('preserves the live workspace when restored ownership cannot be applied', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_chown_failure' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'preserve ownership-failure state'); + session!.markPrimed('stale.csv', 'stale-cache-key'); + const ownershipFailure = Object.assign(new Error('ownership denied'), { code: 'EPERM' }); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive({ 'restored.csv': 'new bytes' })) as never, + res as never, + { + allowUnprivilegedOwnershipFallback: false, + ownershipOps: { + lchown: async () => { throw ownershipFailure; }, + chown: async () => { throw ownershipFailure; }, + }, + }, + ); + + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')) + .toBe('preserve ownership-failure state'); + expect(session!.primedInputId('stale.csv')).toBe('stale-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('a corrupt archive fails without touching the live workspace', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_2' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + await fsp.writeFile(path.join(dir, 'live.txt'), 'preserve corrupt-archive state'); + session!.markPrimed('live.txt', 'live-cache-key'); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(Buffer.from('not a tarball')) as never, + res as never, + ); + + expect(res.statusCode).toBe(500); + expect(await fsp.readFile(path.join(dir, 'live.txt'), 'utf8')) + .toBe('preserve corrupt-archive state'); + expect(session!.primedInputId('live.txt')).toBe('live-cache-key'); + expect(session!.dirtyReason).toBeUndefined(); + }); + + test('does not load or surface metadata from the incompatible v1 identity schema', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_v1' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + const v1 = JSON.stringify({ + marker: 'codeapi.session-meta.v1', + primed: [['input.csv', { + id: 'per-execution-masked-id', + sessionId: 'per-execution-masked-session', + readOnly: false, + }]], + surfaced: [], + }); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive({ + 'input.csv': 'sandbox-modified', + [SESSION_META_FILE]: v1, + })) as never, + res as never, + ); + + expect(res.statusCode).toBe(200); + expect(session!.primedInputId('input.csv')).toBeUndefined(); + expect(await fsp.lstat(path.join(dir, SESSION_META_FILE)).catch(() => null)).toBeNull(); + }); + + test('restores metadata from a compatible legacy in-workspace sidecar', async () => { + config.session_workspace_enabled = true; + const session = bindSessionWorkspace({ runtimeSessionId: 'rt_restore_legacy_v2' }); + seedNonRootIdentity(session!); + const { dir } = await session!.ownership(); + const v2 = JSON.stringify({ + marker: SESSION_META_MARKER, + primed: [['input.csv', { + id: 'stable-cache-key', + readOnly: false, + hash: 'original-hash', + }]], + surfaced: [['output.csv', 'output-hash']], + }); + + const res = fakeStreamRes(); + await restoreSessionCheckpoint( + Readable.from(await makeArchive({ + 'input.csv': 'sandbox-modified', + [SESSION_META_FILE]: v2, + })) as never, + res as never, + ); + + expect(res.statusCode).toBe(200); + expect(session!.primedInputId('input.csv')).toBe('stable-cache-key'); + expect(session!.isSurfaced('output.csv', 'output-hash')).toBe(true); + expect(await fsp.lstat(path.join(dir, SESSION_META_FILE)).catch(() => null)).toBeNull(); + }); +}); diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts new file mode 100644 index 0000000..8ddd739 --- /dev/null +++ b/api/src/session-checkpoint.ts @@ -0,0 +1,533 @@ +import { spawn } from 'child_process'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import type { Request, Response } from 'express'; +import { Transform } from 'stream'; +import { pipeline } from 'stream/promises'; +import { createGunzip, createGzip } from 'zlib'; +import { config } from './config'; +import { logger } from './logger'; +import { SANDBOX_WORKSPACE_ROOT, SESSION_WORKSPACE_ID } from './workspace-isolation'; +import type { SessionMetaSnapshot, SessionWorkspace } from './session-workspace'; +import { SESSION_META_FILE, SESSION_META_MARKER, getBoundSessionWorkspace } from './session-workspace'; + +/** + * Session workspace checkpoint / restore. + * + * Makes an expiring MicroVM's state survive across a relaunch: the control + * plane pulls a compressed archive of the session workspace over the authed + * proxy (GET /checkpoint), stores it in S3, and pushes it back into a fresh + * VM's workspace before the first execute (POST /restore). The untrusted VM + * never touches S3 — only tars its own `/mnt/data`. + * + * Only reachable when a session is bound (getBoundSessionWorkspace); returns + * 409 otherwise so the legacy fresh-per-job runner exposes nothing new. + */ + +const CHECKPOINT_CONTENT_TYPE = 'application/x-gtar'; +/** Runner-only top-level archive member. It is a sibling of `session/`, never + * a path inside the user workspace, so user code cannot collide with it. Keep + * this as ONE top-level file: pre-control-namespace restores use + * `--strip-components=1`, which safely skips a one-component member but would + * extract `control-dir/file` as a user-visible `file` during rollback. */ +const CHECKPOINT_CONTROL_FILE = '.codeapi-checkpoint-control.v2.json'; +/** Primed-input state is correctness-critical: losing it can overwrite + * sandbox-modified inputs after restore. Keep the full control record whenever + * it fits this bounded parse budget. Only beyond the hard limit may expendable + * surfaced-output signatures be dropped; primed state is never discarded. */ +const CHECKPOINT_CONTROL_MAX_BYTES = 16 * 1024 * 1024; + +export class SessionCheckpointError extends Error {} + +function currentUid(): number | undefined { + return typeof process.getuid === 'function' ? process.getuid() : undefined; +} + +/** + * Serializes best-effort control metadata without making an otherwise valid + * workspace impossible to checkpoint. Surfaced-output signatures are only a + * dedup optimization, so discard them first. Primed-input state protects + * sandbox edits from being overwritten after restore, so it is never silently + * discarded. An extreme session that exceeds the bounded metadata budget must + * fail its checkpoint instead of persisting a corrupt recovery point. + */ +function checkpointControlBytes(session: SessionWorkspace): Buffer { + const snapshot = session.snapshotMeta(); + const encode = (meta: SessionMetaSnapshot): Buffer => + Buffer.from(JSON.stringify({ marker: SESSION_META_MARKER, ...meta })); + + const full = encode(snapshot); + if (full.length <= CHECKPOINT_CONTROL_MAX_BYTES) return full; + + const withoutSurfaced = encode({ primed: snapshot.primed, surfaced: [] }); + logger.warn( + { + fullBytes: full.length, + reducedBytes: withoutSurfaced.length, + maxBytes: CHECKPOINT_CONTROL_MAX_BYTES, + }, + 'Checkpoint control metadata is oversized; dropping surfaced-output state', + ); + if (withoutSurfaced.length <= CHECKPOINT_CONTROL_MAX_BYTES) { + return withoutSurfaced; + } + + throw new SessionCheckpointError( + `checkpoint primed-input metadata exceeds ${CHECKPOINT_CONTROL_MAX_BYTES} bytes`, + ); +} + +export interface SessionCheckpointStreamOptions { + /** Test seam for deterministic late tar-failure coverage. */ + tarCommand?: string; +} + +/** Streams a capped `tar.gz` of the session workspace to the response. */ +export async function streamSessionCheckpoint( + res: Response, + options: SessionCheckpointStreamOptions = {}, +): Promise { + const session = getBoundSessionWorkspace(); + if (!session) { + res.status(409).json({ message: 'No session workspace is bound' }); + return; + } + if (session.dirtyReason) { + res.status(409).json({ + error: 'session_workspace_dirty', + message: 'Session workspace must be restored before checkpointing', + }); + return; + } + await session.ownership(); + + /* Carry priming/output-diff state in a separate top-level archive member. + * The old format temporarily wrote SESSION_META_FILE into `/mnt/data`, which + * either deleted a symlink/directory with that user-visible name or skipped + * metadata for a legitimate regular file. A sibling archive member has no + * collision with any path user code can create and never mutates the live + * workspace. Restore still understands the legacy in-workspace sidecar. */ + let controlStage: string | undefined; + try { + controlStage = await fsp.mkdtemp(path.join(os.tmpdir(), 'codeapi-checkpoint-control-')); + const controlBytes = checkpointControlBytes(session); + await fsp.writeFile( + path.join(controlStage, CHECKPOINT_CONTROL_FILE), + controlBytes, + { flag: 'wx', mode: 0o600 }, + ); + + res.status(200); + res.setHeader('Content-Type', CHECKPOINT_CONTENT_TYPE); + /* Cap the uncompressed tar before gzip as well as the compressed response. + * Restore applies the same two limits in reverse, so every checkpoint that + * creation completes is guaranteed to be admissible under the same config. + * Keeping tar and gzip as separate streaming stages preserves bounded + * memory, backpressure, and the existing tar.gz wire format. */ + const tar = spawn(options.tarCommand ?? 'tar', [ + '-cf', '-', + '-C', SANDBOX_WORKSPACE_ROOT, SESSION_WORKSPACE_ID, + '-C', controlStage, CHECKPOINT_CONTROL_FILE, + ], { + stdio: ['ignore', 'pipe', 'pipe'], + env: { ...process.env, COPYFILE_DISABLE: '1' }, + }); + tar.stderr.on('data', (chunk: Buffer) => logger.debug({ tar: chunk.toString() }, 'checkpoint tar')); + /* Register the 'close' listener BEFORE awaiting the pipeline: for a small + * workspace tar can exit and emit 'close' before pipeline resolves, and a + * listener attached only afterward would miss it and hang here forever. + * The 'error' listener turns a spawn failure (e.g. tar missing from PATH) + * into a rejected promise instead of an unhandled ChildProcess 'error' + * crashing the runner. */ + const closed: Promise = new Promise((resolve, reject) => { + tar.on('close', resolve); + tar.on('error', reject); + }); + /* Observe the rejection immediately: when `tar` fails to spawn, `pipeline` + * below rejects first and we never reach `await closed` — an unobserved + * rejection would then take the whole runner down after we already + * answered 500. The later `await closed` still sees the same rejection. */ + closed.catch(() => {}); + try { + await pipeline( + tar.stdout, + checkpointStreamLimit(config.checkpoint_max_bytes, 'expanded'), + createGzip(), + checkpointStreamLimit(config.checkpoint_max_bytes, 'compressed'), + res, + /* Do not publish a successful HTTP EOF until tar itself exits zero. + * A tar process can emit bytes and close stdout before reporting a + * late read error. Ending the response inside pipeline would make that + * partial archive look complete to the control plane, which could then + * commit it and prune the preceding recovery point. */ + { end: false }, + ); + const code = await closed; + if (code !== 0) throw new SessionCheckpointError(`checkpoint tar exited ${code}`); + res.end(); + } catch (error) { + /* A downstream disconnect or either byte cap can stop consumption while + * tar is still writing. Terminate and reap it so no child or temp-stage + * lifetime escapes the failed request. */ + if (tar.exitCode === null && tar.signalCode === null) tar.kill('SIGKILL'); + await closed.catch(() => {}); + throw error; + } + } catch (error) { + logger.error({ err: error }, 'Failed to stream session checkpoint'); + if (!res.headersSent) res.status(500).json({ message: 'checkpoint failed' }); + else res.destroy(); + } finally { + if (controlStage) { + await fsp.rm(controlStage, { recursive: true, force: true }).catch(() => {}); + } + } +} + +export interface SessionCheckpointOwnershipOps { + lchown(target: string, uid: number, gid: number): Promise; + chown(target: string, uid: number, gid: number): Promise; +} + +export interface SessionCheckpointRestoreOptions { + /** Test seam for deterministic ownership-failure coverage. */ + ownershipOps?: SessionCheckpointOwnershipOps; + /** Test seam for deterministic commit/rollback failure coverage. */ + rename?: (source: string, destination: string) => Promise; + /** Non-root local development cannot chown to the compatibility sandbox UID. + * Production hardened/per-job-UID mode never enables this fallback. */ + allowUnprivilegedOwnershipFallback?: boolean; +} + +const DEFAULT_OWNERSHIP_OPS: SessionCheckpointOwnershipOps = { + lchown: (target, uid, gid) => fsp.lchown(target, uid, gid), + chown: (target, uid, gid) => fsp.chown(target, uid, gid), +}; + +function errorCode(error: unknown): string | undefined { + return typeof error === 'object' && error !== null && 'code' in error + ? String((error as { code?: unknown }).code) + : undefined; +} + +function checkpointStreamLimit(maxBytes: number, kind: 'compressed' | 'expanded'): Transform { + let received = 0; + return new Transform({ + transform(chunk: Buffer, _encoding, callback) { + received += chunk.length; + callback( + received > maxBytes + ? new SessionCheckpointError(`checkpoint ${kind} stream exceeds ${maxBytes} bytes`) + : null, + chunk, + ); + }, + }); +} + +/** Extracts a `tar.gz` into a staging tree, validates its allowed top-level + * members, re-owns the workspace, then replaces the live session directory. + * Staging keeps a corrupt archive or ownership failure from ever becoming an + * executable workspace. */ +export async function restoreSessionCheckpoint( + req: Request, + res: Response, + options: SessionCheckpointRestoreOptions = {}, +): Promise { + const session = getBoundSessionWorkspace(); + if (!session) { + res.status(409).json({ message: 'No session workspace is bound' }); + return; + } + const { dir, uid, gid } = await session.ownership(); + const restoreStage = await fsp.mkdtemp( + path.join(SANDBOX_WORKSPACE_ROOT, '.session-restore-'), + ); + const restoredWorkspace = path.join(restoreStage, SESSION_WORKSPACE_ID); + const restoredControl = path.join(restoreStage, CHECKPOINT_CONTROL_FILE); + const liveBackup = path.join(restoreStage, '.live-workspace-backup'); + const rename = options.rename ?? fsp.rename; + let liveBackedUp = false; + let committed = false; + let retainStageForRecovery = false; + try { + /* Decompress in-process so the runner can independently cap BOTH the + * compressed request and the expanded tar stream. Passing gzip directly to + * `tar -xzf` would let a tiny, highly-compressible checkpoint fill the + * workspace disk before any post-extraction validation could run. */ + const gunzip = createGunzip(); + /* Keep reading through zero end markers until the expanded stream reaches + * EOF. Without --ignore-zeros, tar can exit after the first marker while + * gunzip still has valid tar padding to emit, closing stdin underneath + * pipeline with ERR_STREAM_PREMATURE_CLOSE. */ + const tar = spawn('tar', ['--ignore-zeros', '-xf', '-', '-C', restoreStage], { + stdio: ['pipe', 'ignore', 'pipe'], + env: { ...process.env, COPYFILE_DISABLE: '1' }, + }); + tar.stderr.on('data', (chunk: Buffer) => logger.debug({ tar: chunk.toString() }, 'restore tar')); + const closed: Promise = new Promise((resolve, reject) => { + tar.on('close', resolve); + tar.on('error', reject); + }); + /* Observe immediately: a spawn failure may reject the input pipeline first. */ + closed.catch(() => {}); + try { + /* Register the 'close' listener before awaiting the pipeline (see the + * create side): a small upload can finish and 'close' can fire before + * pipeline resolves, and a listener attached afterward would hang. */ + await pipeline( + req, + checkpointStreamLimit(config.checkpoint_max_bytes, 'compressed'), + gunzip, + checkpointStreamLimit(config.checkpoint_max_bytes, 'expanded'), + tar.stdin, + ); + const code = await closed; + if (code !== 0) throw new SessionCheckpointError(`restore tar exited ${code}`); + } catch (error) { + if (tar.exitCode === null && tar.signalCode === null) tar.kill('SIGKILL'); + await closed.catch(() => {}); + throw error; + } + + const topLevel = await fsp.readdir(restoreStage, { withFileTypes: true }); + const allowed = new Set([SESSION_WORKSPACE_ID, CHECKPOINT_CONTROL_FILE]); + const unexpected = topLevel.find(entry => !allowed.has(entry.name)); + if (unexpected) { + throw new SessionCheckpointError( + `checkpoint contains unexpected top-level member: ${unexpected.name}`, + ); + } + const workspaceStat = await fsp.lstat(restoredWorkspace).catch(() => null); + if (!workspaceStat?.isDirectory() || workspaceStat.isSymbolicLink()) { + throw new SessionCheckpointError('checkpoint session workspace is not a real directory'); + } + const controlStat = await fsp.lstat(restoredControl).catch(() => null); + if (controlStat && (!controlStat.isFile() || controlStat.isSymbolicLink())) { + throw new SessionCheckpointError('checkpoint control metadata is not a regular file'); + } + + const restoredMeta = await readRestoredMeta( + restoredWorkspace, + controlStat ? restoredControl : undefined, + ); + const runnerUid = currentUid(); + const allowUnprivilegedFallback = + options.allowUnprivilegedOwnershipFallback + ?? ( + runnerUid !== undefined + && runnerUid !== 0 + && !config.hardened_sandbox_mode + && !config.per_job_uids + ); + await chownRecursive( + restoredWorkspace, + uid, + gid, + options.ownershipOps ?? DEFAULT_OWNERSHIP_OPS, + allowUnprivilegedFallback, + ); + + /* Commit only after extraction, metadata validation and ownership all + * succeeded. Rename the old workspace aside first so a failed install can + * atomically put it back; deleting it before rename made a staged failure + * destroy valid live session state. Both paths share the workspace + * filesystem, so each individual rename is atomic. */ + await rename(dir, liveBackup); + liveBackedUp = true; + await rename(restoredWorkspace, dir); + liveBackedUp = false; + session.loadMeta(restoredMeta); + committed = true; + + /* The checkpoint is committed before delivery of the HTTP acknowledgement. + * A socket/serialization error now causes the control plane to recycle this + * VM, but must never roll the successfully restored workspace backward. */ + await fsp.rm(liveBackup, { recursive: true, force: true }).catch(error => { + logger.warn({ err: error, liveBackup }, 'Failed to remove replaced checkpoint workspace'); + }); + res.status(200).json({ status: 'restored', dir: path.basename(dir) }); + } catch (error) { + if (committed) { + logger.error({ err: error }, 'Checkpoint restored but response delivery failed'); + throw error; + } + + if (liveBackedUp) { + try { + await rename(liveBackup, dir); + liveBackedUp = false; + } catch (rollbackError) { + retainStageForRecovery = true; + session.markDirty('checkpoint restore rollback failed'); + logger.error( + { err: error, rollbackErr: rollbackError, recoveryPath: liveBackup }, + 'Failed to roll back checkpoint workspace replacement', + ); + } + } + logger.error({ err: error }, 'Failed to restore session checkpoint'); + /* Extraction, validation, and ownership happen only in restoreStage. + * Ordinary pre-commit failures therefore leave both the live workspace and + * its matching metadata/dirty state untouched. Only an unsuccessful + * rollback above marks the session dirty and retains its recovery copy. */ + if (!res.headersSent) res.status(500).json({ message: 'restore failed' }); + } finally { + if (!retainStageForRecovery) { + await fsp.rm(restoreStage, { recursive: true, force: true }).catch(() => {}); + } + } +} + +function isSessionMetaSnapshot(value: unknown): value is SessionMetaSnapshot { + if (value == null || typeof value !== 'object') return false; + const snapshot = value as SessionMetaSnapshot; + if (snapshot.marker !== SESSION_META_MARKER) return false; + if (!Array.isArray(snapshot.primed) || !Array.isArray(snapshot.surfaced)) return false; + const validPrimed = snapshot.primed.every(entry => + Array.isArray(entry) + && entry.length === 2 + && typeof entry[0] === 'string' + && entry[1] != null + && typeof entry[1] === 'object' + && typeof entry[1].id === 'string' + && typeof entry[1].readOnly === 'boolean' + && (entry[1].hash === undefined || typeof entry[1].hash === 'string')); + const validSurfaced = snapshot.surfaced.every(entry => + Array.isArray(entry) + && entry.length === 2 + && typeof entry[0] === 'string' + && typeof entry[1] === 'string'); + return validPrimed && validSurfaced; +} + +/** Applies runner control metadata. New checkpoints carry it outside + * `session/`; old checkpoints are read from the legacy in-workspace sidecar. + * Presence of a new control member is authoritative: invalid metadata rejects + * the restore, and a user file at the legacy name is never reinterpreted or + * deleted. */ +async function readRestoredMeta( + workspaceDir: string, + controlPath?: string, +): Promise { + const empty: SessionMetaSnapshot = { primed: [], surfaced: [] }; + /* Parse and validate metadata without mutating the live session. The caller + * installs the returned snapshot only after the staged workspace commits. */ + if (controlPath) { + const stat = await fsp.lstat(controlPath).catch(() => null); + if (!stat?.isFile() || stat.isSymbolicLink()) { + throw new SessionCheckpointError('checkpoint control metadata is not a regular file'); + } + if (stat.size > CHECKPOINT_CONTROL_MAX_BYTES) { + throw new SessionCheckpointError( + `checkpoint control metadata exceeds ${CHECKPOINT_CONTROL_MAX_BYTES} bytes`, + ); + } + let parsed: unknown; + try { + parsed = JSON.parse(await fsp.readFile(controlPath, 'utf8')) as unknown; + } catch { + throw new SessionCheckpointError('checkpoint control metadata is not valid JSON'); + } + if (!isSessionMetaSnapshot(parsed)) { + throw new SessionCheckpointError('checkpoint control metadata is malformed or incompatible'); + } + return parsed; + } + + const metaPath = path.join(workspaceDir, SESSION_META_FILE); + try { + /* Backward compatibility for checkpoints created by the pre-control- + * namespace format. Only trust a runner-owned writable regular file: + * user-created and read-only-input collisions remain ordinary user data. */ + const stat = await fsp.lstat(metaPath).catch(() => null); + if (!stat?.isFile()) return empty; + const trustedOwner = currentUid(); + const ownerMatches = trustedOwner == null || stat.uid === trustedOwner; + const ownerWritable = (stat.mode & 0o200) !== 0; + if (!ownerMatches || !ownerWritable) { + logger.warn('Ignoring untrusted session meta sidecar from restored workspace'); + return empty; + } + const parsed = JSON.parse(await fsp.readFile(metaPath, 'utf8')) as unknown; + if (isSessionMetaSnapshot(parsed)) { + await fsp.rm(metaPath, { force: true }).catch(() => {}); + return parsed; + } else if ( + typeof (parsed as { marker?: unknown })?.marker === 'string' && + (parsed as { marker: string }).marker.startsWith('codeapi.session-meta.v') + ) { + /* This is a runner-owned sidecar from an incompatible pre-release image, + * not a user file. Never load it under the current schema and never leave + * it in the workspace to be surfaced as a generated artifact. Rollouts + * that change this marker must drain/recycle old development sessions. */ + logger.warn( + { marker: (parsed as { marker: string }).marker, expected: SESSION_META_MARKER }, + 'Ignoring incompatible session checkpoint metadata', + ); + await fsp.rm(metaPath, { force: true }).catch(() => {}); + } + } catch (error) { + logger.debug({ err: error }, 'No session meta sidecar to restore'); + } + return empty; +} + +async function applyOwnership( + operation: () => Promise, + allowUnprivilegedFallback: boolean, +): Promise { + try { + await operation(); + } catch (error) { + const code = errorCode(error); + if (allowUnprivilegedFallback && (code === 'EPERM' || code === 'EACCES')) { + return; + } + throw error; + } +} + +async function chownRecursive( + dir: string, + uid: number, + gid: number, + ownershipOps: SessionCheckpointOwnershipOps, + allowUnprivilegedFallback: boolean, +): Promise { + await applyOwnership( + () => ownershipOps.lchown(dir, uid, gid), + allowUnprivilegedFallback, + ); + const entries = await fsp.readdir(dir, { withFileTypes: true }); + for (const entry of entries) { + const full = path.join(dir, entry.name); + /* A restored checkpoint is untrusted content: never follow symlinks. A + * `session/x -> /etc/passwd` entry would otherwise have `chown` re-own the + * target outside the workspace. Do not trust Dirent.d_type here because + * some filesystems report DT_UNKNOWN; lstat every member, lchown links, and + * never recurse through them. */ + const stat = await fsp.lstat(full); + if (stat.isSymbolicLink()) { + await applyOwnership( + () => ownershipOps.lchown(full, uid, gid), + allowUnprivilegedFallback, + ); + continue; + } + await applyOwnership( + () => ownershipOps.chown(full, uid, gid), + allowUnprivilegedFallback, + ); + if (stat.isDirectory()) { + await chownRecursive( + full, + uid, + gid, + ownershipOps, + allowUnprivilegedFallback, + ); + } + } +} diff --git a/api/src/session-inputs.prime.test.ts b/api/src/session-inputs.prime.test.ts new file mode 100644 index 0000000..5fe5d2a --- /dev/null +++ b/api/src/session-inputs.prime.test.ts @@ -0,0 +1,200 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { Job, type TFile } from './job'; +import { SESSION_INPUT_CACHE_DIR, inputCacheKey } from './session-inputs'; +import { SessionWorkspace } from './session-workspace'; +import { fallbackSandboxIdentity } from './workspace-isolation'; + +/** + * The seam that makes the redesign work: a pushed cache entry must prime + * exactly as a file-server fetch would. Unit tests cover the cache, and route + * tests cover the wiring — this covers the join, which is where a live-only + * failure hid. + */ + +let tmpDir: string; + +afterEach(async () => { + await fsp.rm(SESSION_INPUT_CACHE_DIR, { recursive: true, force: true }).catch(() => {}); + if (tmpDir) await fsp.rm(tmpDir, { recursive: true, force: true }).catch(() => {}); +}); + +async function seedCache(sid: string, id: string, body: string, meta: object): Promise { + await fsp.mkdir(SESSION_INPUT_CACHE_DIR, { recursive: true }); + const key = inputCacheKey(sid, id); + await fsp.writeFile(path.join(SESSION_INPUT_CACHE_DIR, key), body); + await fsp.writeFile(path.join(SESSION_INPUT_CACHE_DIR, `${key}.json`), JSON.stringify(meta)); +} + +function makeJob(files: TFile[], session?: object): Job { + return new Job({ + session_id: 'prime-cache-test', + runtime: { language: 'bash', version: '5.0.0', aliases: [], runtime: 'bash' } as never, + args: [], + stdin: '', + files, + timeouts: { run: 5000, compile: 5000 }, + cpu_times: { run: 5000, compile: 5000 }, + memory_limits: { run: 128 * 1024 * 1024, compile: 128 * 1024 * 1024 }, + session, + } as never); +} + +describe('priming from the pushed input cache', () => { + test('writes the cached bytes under the requested name without any HTTP fetch', async () => { + tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-')); + await seedCache('s1', 'f1', 'a,b\n1,2\n', { readOnly: false }); + + const file: TFile = { id: 'f1', storage_session_id: 's1', name: 'data.csv' }; + const job = makeJob([file]); + (job as unknown as { submissionDir: string }).submissionDir = tmpDir; + + /* A fetch would fail here: no file server is configured in tests, so a + * successful write proves the bytes came from the cache. */ + const written = await job.downloadAndWriteFile(file); + expect(written).toBe('data.csv'); + expect(await fsp.readFile(path.join(tmpDir, 'data.csv'), 'utf8')).toBe('a,b\n1,2\n'); + }); + + test('honors the cached read-only bit at the requested destination', async () => { + tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-ro-')); + await seedCache('s1', 'ro', 'SKILL\n', { readOnly: true }); + + const file: TFile = { id: 'ro', storage_session_id: 's1', name: 'requested.md' }; + const job = makeJob([file]); + (job as unknown as { submissionDir: string }).submissionDir = tmpDir; + + const written = await job.downloadAndWriteFile(file); + expect(written).toBe('requested.md'); + const stat = await fsp.lstat(path.join(tmpDir, 'requested.md')); + expect(stat.mode & 0o777).toBe(0o444); + }); + + test('one cached object serves several destinations in the same execute', async () => { + tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-multi-')); + await seedCache('s1', 'shared', 'shared bytes\n', { readOnly: false }); + + /* Live regression: a per-OBJECT cached name made the second ref resolve to + * the first ref's path — which overwrote a file the sandbox had edited. + * The cache is keyed by object, so destinations belong to the refs. */ + const job = makeJob([]); + (job as unknown as { submissionDir: string }).submissionDir = tmpDir; + for (const name of ['first.txt', 'nested/second.txt']) { + const written = await job.downloadAndWriteFile({ id: 'shared', storage_session_id: 's1', name }); + expect(written).toBe(name); + } + expect(await fsp.readFile(path.join(tmpDir, 'first.txt'), 'utf8')).toBe('shared bytes\n'); + expect(await fsp.readFile(path.join(tmpDir, 'nested/second.txt'), 'utf8')).toBe('shared bytes\n'); + }); + + test('stable cache identity preserves a modified session input across fresh grant handles', async () => { + tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-stable-')); + await fsp.writeFile(path.join(tmpDir, 'data.csv'), 'sandbox-modified\n'); + const stable = inputCacheKey('raw-session', 'raw-object'); + const session = { + primedInputId: () => stable, + primedHash: () => 'original-upload-hash', + isPrimedReadOnly: () => false, + }; + const job = makeJob([], session); + const internals = job as unknown as { + submissionDir: string; + reusePrimedInput(file: TFile): Promise; + inputFileHashes: Map; + }; + internals.submissionDir = tmpDir; + + const reused = await internals.reusePrimedInput({ + id: 'second-grant-object-handle', + storage_session_id: 'second-grant-session-handle', + input_cache_key: stable, + name: 'data.csv', + }); + + expect(reused).toBe(true); + expect(await fsp.readFile(path.join(tmpDir, 'data.csv'), 'utf8')).toBe('sandbox-modified\n'); + expect(internals.inputFileHashes.get('data.csv')).toMatchObject({ + originalId: 'second-grant-object-handle', + originalSessionId: 'second-grant-session-handle', + hash: 'original-upload-hash', + }); + }); + + test('same-session re-prime refuses a symlinked ancestor without touching its target', async () => { + tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-ancestor-link-')); + const outsideDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-outside-')); + const outsideTarget = path.join(outsideDir, 'input.txt'); + await fsp.writeFile(outsideTarget, 'outside-original\n'); + await fsp.symlink(outsideDir, path.join(tmpDir, 'nested')); + await seedCache('s1', 'f1', 'delivered\n', { readOnly: false }); + + const file: TFile = { + id: 'f1', + storage_session_id: 's1', + name: 'nested/input.txt', + }; + const session = new SessionWorkspace({ runtimeSessionId: 'rt_symlink_ancestor' }); + session.markPrimed(file.name, inputCacheKey('s1', 'f1')); + const job = makeJob([file], session); + const primeInputFile = ( + job as unknown as { + primeInputFile( + input: TFile, + context: { submissionDir: string; identity: ReturnType }, + ): Promise; + } + ).primeInputFile.bind(job); + + try { + await expect(primeInputFile(file, { + submissionDir: tmpDir, + identity: fallbackSandboxIdentity(), + })).rejects.toThrow('symlinked workspace path'); + expect(await fsp.readFile(outsideTarget, 'utf8')).toBe('outside-original\n'); + } finally { + await fsp.rm(outsideDir, { recursive: true, force: true }); + } + }); + + test('same-session re-prime replaces a direct symlink instead of following it', async () => { + tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-target-link-')); + const outsideDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'prime-cache-outside-')); + const outsideTarget = path.join(outsideDir, 'outside.txt'); + const destination = path.join(tmpDir, 'input.txt'); + await fsp.writeFile(outsideTarget, 'outside-original\n'); + await fsp.symlink(outsideTarget, destination); + await seedCache('s1', 'f1', 'delivered\n', { readOnly: false }); + + const file: TFile = { id: 'f1', storage_session_id: 's1', name: 'input.txt' }; + const session = new SessionWorkspace({ runtimeSessionId: 'rt_symlink_target' }); + session.markPrimed(file.name, inputCacheKey('s1', 'f1')); + const job = makeJob([file], session); + const primeInputFile = ( + job as unknown as { + primeInputFile( + input: TFile, + context: { submissionDir: string; identity: ReturnType }, + ): Promise; + } + ).primeInputFile.bind(job); + + try { + await primeInputFile(file, { + submissionDir: tmpDir, + identity: fallbackSandboxIdentity(), + }); + + expect((await fsp.lstat(destination)).isFile()).toBe(true); + expect(await fsp.readFile(destination, 'utf8')).toBe('delivered\n'); + expect(await fsp.readFile(outsideTarget, 'utf8')).toBe('outside-original\n'); + } finally { + await fsp.rm(outsideDir, { recursive: true, force: true }); + } + }); +}); diff --git a/api/src/session-inputs.test.ts b/api/src/session-inputs.test.ts new file mode 100644 index 0000000..639ccad --- /dev/null +++ b/api/src/session-inputs.test.ts @@ -0,0 +1,322 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { spawnSync } from 'child_process'; +import { Readable } from 'stream'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { gzipSync, gunzipSync } from 'zlib'; +import { SANDBOX_WORKSPACE_ROOT, reapStaleWorkspaces } from './workspace-isolation'; +import { + SESSION_INPUT_CACHE_DIR, + cachedInputResponse, + hasCachedInput, + inputCacheKey, + openCachedInput, + pruneInputCache, + resolveSessionInputCacheDir, + storeCachedInputs, + writeAllToHandle, +} from './session-inputs'; + +afterEach(async () => { + await fsp.rm(SESSION_INPUT_CACHE_DIR, { recursive: true, force: true }).catch(() => {}); +}); + +/** Builds the digest-named batch the control plane pushes. */ +async function makeBatch( + entries: Array<{ storageSessionId: string; id: string; body: string; meta?: object }>, + extras: Record = {}, +): Promise { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'inputs-batch-')); + for (const entry of entries) { + const key = inputCacheKey(entry.storageSessionId, entry.id); + await fsp.writeFile(path.join(tmp, key), entry.body); + if (entry.meta) await fsp.writeFile(path.join(tmp, `${key}.json`), JSON.stringify(entry.meta)); + } + for (const [name, body] of Object.entries(extras)) { + await fsp.mkdir(path.dirname(path.join(tmp, name)), { recursive: true }); + await fsp.writeFile(path.join(tmp, name), body); + } + const tar = spawnSync('tar', ['-czf', '-', '-C', tmp, '.'], { + maxBuffer: 64 * 1024 * 1024, + env: { ...process.env, COPYFILE_DISABLE: '1' }, + }); + await fsp.rm(tmp, { recursive: true, force: true }); + if (tar.status !== 0) throw new Error(`fixture tar exited ${tar.status}`); + return tar.stdout; +} + +/** + * Cross-component contract. The control plane names batch members with ITS + * implementation of this digest and the runner looks entries up with THIS one; + * nothing else ties them together, so both suites assert the same vector. + * + * They silently diverged once — one side used a literal NUL separator, the + * other a space — and every push landed in the cache under names no lookup + * would ever produce. Both unit suites passed (each self-consistent) and only + * a live execution showed it, as "file not found" for files that had been + * delivered successfully. + */ +const GOLDEN_KEY_SID_1_FILE_1 = 'a995f1e7977466c5636419d21582e0b44420c44d2d7e2660b13aa4d4b4667d90'; + +describe('pushed input cache', () => { + test('digests match the wire contract the control plane names members with', () => { + expect(inputCacheKey('sid-1', 'file-1')).toBe(GOLDEN_KEY_SID_1_FILE_1); + }); + + test('retries short FileHandle writes until the complete slice is persisted', async () => { + const written: Buffer[] = []; + const handle = { + write: async ( + buffer: Buffer, + offset: number, + length: number, + ): Promise<{ bytesWritten: number; buffer: Buffer }> => { + const bytesWritten = Math.max(1, Math.floor(length / 2)); + written.push(Buffer.from(buffer.subarray(offset, offset + bytesWritten))); + return { bytesWritten, buffer }; + }, + }; + await writeAllToHandle(handle as never, Buffer.from('complete-payload')); + expect(Buffer.concat(written).toString()).toBe('complete-payload'); + expect(written.length).toBeGreaterThan(1); + }); + + test('recovers from ENOSPC mid-write without duplicating the persisted prefix', async () => { + const written: Buffer[] = []; + let calls = 0; + let recoveries = 0; + const handle = { + write: async ( + buffer: Buffer, + offset: number, + length: number, + ): Promise<{ bytesWritten: number; buffer: Buffer }> => { + calls += 1; + if (calls === 2) { + throw Object.assign(new Error('disk full'), { code: 'ENOSPC' }); + } + const bytesWritten = calls === 1 ? Math.min(4, length) : length; + written.push(Buffer.from(buffer.subarray(offset, offset + bytesWritten))); + return { bytesWritten, buffer }; + }, + }; + + await writeAllToHandle( + handle as never, + Buffer.from('complete-payload'), + async () => { recoveries += 1; }, + ); + + expect(recoveries).toBe(1); + expect(Buffer.concat(written).toString()).toBe('complete-payload'); + }); + + test('lives outside the workspace root so the reaper cannot eat it', async () => { + /* Regression: the cache was originally a dot-directory INSIDE + * SANDBOX_WORKSPACE_ROOT, where the stale-workspace reaper treats every + * entry as a workspace — it deleted pending inputs between the push and + * the execute they were pushed for. */ + expect(SESSION_INPUT_CACHE_DIR.startsWith(`${SANDBOX_WORKSPACE_ROOT}/`)).toBe(false); + + await fsp.mkdir(SESSION_INPUT_CACHE_DIR, { recursive: true }); + const key = inputCacheKey('s1', 'survives'); + await fsp.writeFile(path.join(SESSION_INPUT_CACHE_DIR, key), 'bytes'); + await fsp.writeFile( + path.join(SESSION_INPUT_CACHE_DIR, `${key}.json`), + JSON.stringify({ readOnly: false }), + ); + await reapStaleWorkspaces({ maxAgeMs: 0 }); + expect(await hasCachedInput('s1', 'survives')).toBe(true); + }); + + test('stores a batch and serves it as the response a fetch would have returned', async () => { + const batch = await makeBatch([ + { storageSessionId: 's1', id: 'f1', body: 'a,b\n1,2\n', meta: { readOnly: false } }, + { storageSessionId: 's1', id: 'ro', body: 'SKILL\n', meta: { readOnly: true } }, + ]); + expect(await storeCachedInputs(Readable.from(batch))).toBe(2); + + expect(await hasCachedInput('s1', 'f1')).toBe(true); + expect(await hasCachedInput('s1', 'nope')).toBe(false); + + const hit = await openCachedInput('s1', 'f1'); + expect(hit).not.toBeNull(); + const response = cachedInputResponse(hit!); + /* No Content-Disposition: the cache is keyed by OBJECT, so the requesting + * ref owns the destination name and priming falls back to it. */ + expect(response.headers.get('content-disposition')).toBeNull(); + expect(response.headers.get('x-read-only')).toBeNull(); + expect(await response.text()).toBe('a,b\n1,2\n'); + + const readOnly = cachedInputResponse((await openCachedInput('s1', 'ro'))!); + expect(readOnly.headers.get('x-read-only')).toBe('true'); + }); + + test('an opened cache hit remains readable when concurrent pruning unlinks its path', async () => { + const batch = await makeBatch([ + { storageSessionId: 's1', id: 'held', body: 'held-open-bytes', meta: { readOnly: false } }, + ]); + await storeCachedInputs(Readable.from(batch)); + const hit = await openCachedInput('s1', 'held'); + expect(hit).not.toBeNull(); + + await pruneInputCache(1); + expect(await fsp.lstat(hit!.path).catch(() => null)).toBeNull(); + expect(await cachedInputResponse(hit!).text()).toBe('held-open-bytes'); + }); + + test('rejects a batch containing anything but digest-named flat members', async () => { + /* Names are runner-computed digests, so a traversal attempt is not a path + * problem to solve — it is simply not a legal member name. Any member that + * is not `<64 hex>[.json]` fails the whole batch, and nothing lands. */ + const cases: Array> = [ + { 'notadigest.txt': 'nope' }, + { 'nested/deep.txt': 'nope' }, + ]; + for (const extras of cases) { + const batch = await makeBatch([ + { storageSessionId: 's1', id: 'f1', body: 'ok', meta: { readOnly: false } }, + ], extras); + await expect(storeCachedInputs(Readable.from(batch))).rejects.toThrow(); + expect(await hasCachedInput('s1', 'f1')).toBe(false); + } + }); + + test('missing or corrupt metadata is a cache miss, never a read-only downgrade', async () => { + const key = inputCacheKey('s1', 'f1'); + await fsp.mkdir(SESSION_INPUT_CACHE_DIR, { recursive: true }); + await fsp.writeFile(path.join(SESSION_INPUT_CACHE_DIR, key), 'payload'); + expect(await openCachedInput('s1', 'f1')).toBeNull(); + await fsp.writeFile(path.join(SESSION_INPUT_CACHE_DIR, `${key}.json`), '{not json'); + + expect(await openCachedInput('s1', 'f1')).toBeNull(); + expect(await hasCachedInput('s1', 'f1')).toBe(false); + }); + + test('rejects incomplete or invalid object/metadata pairs before committing', async () => { + const missingMeta = await makeBatch([ + { storageSessionId: 's1', id: 'missing', body: 'bytes' }, + ]); + await expect(storeCachedInputs(Readable.from(missingMeta))).rejects.toThrow( + 'requires exactly one data file and metadata sidecar', + ); + + const invalidMeta = await makeBatch([ + { storageSessionId: 's1', id: 'invalid', body: 'bytes', meta: {} }, + ]); + await expect(storeCachedInputs(Readable.from(invalidMeta))).rejects.toThrow( + 'has invalid metadata', + ); + expect(await hasCachedInput('s1', 'missing')).toBe(false); + expect(await hasCachedInput('s1', 'invalid')).toBe(false); + }); + + test('rejects an oversized metadata sidecar before reading it into memory', async () => { + const oversized = await makeBatch([ + { + storageSessionId: 's1', + id: 'oversized-meta', + body: 'bytes', + meta: { readOnly: false, padding: 'x'.repeat(2048) }, + }, + ]); + await expect(storeCachedInputs(Readable.from(oversized))).rejects.toThrow( + 'metadata exceeds', + ); + expect(await hasCachedInput('s1', 'oversized-meta')).toBe(false); + }); + + test('bounds decompressed zero padding after the tar terminator', async () => { + const batch = await makeBatch([ + { + storageSessionId: 's1', + id: 'zero-tail', + body: '', + meta: { readOnly: false }, + }, + ]); + const expanded = gunzipSync(batch); + const zeroTailBomb = gzipSync(Buffer.concat([ + expanded, + Buffer.alloc(2 * 1024 * 1024), + ])); + + await expect( + storeCachedInputs(Readable.from(zeroTailBomb), 1024), + ).rejects.toThrow('Decompressed session input batch exceeds'); + expect(await hasCachedInput('s1', 'zero-tail')).toBe(false); + }); + + test('rejects cache roots that could delete workspaces or broad system paths', () => { + expect(() => resolveSessionInputCacheDir('relative/cache')).toThrow('absolute path'); + expect(() => resolveSessionInputCacheDir('/tmp')).toThrow('dedicated directory'); + expect(() => resolveSessionInputCacheDir('/var')).toThrow('dedicated directory'); + expect(() => resolveSessionInputCacheDir(SANDBOX_WORKSPACE_ROOT)).toThrow('dedicated directory'); + expect(() => resolveSessionInputCacheDir(`${SANDBOX_WORKSPACE_ROOT}/inputs`)).toThrow( + 'dedicated directory', + ); + expect(resolveSessionInputCacheDir('/tmp/sandbox-inputs-safe')).toBe( + '/tmp/sandbox-inputs-safe', + ); + }); + + test('eviction drops least-recently-used entries with their metadata', async () => { + const batch = await makeBatch([ + { storageSessionId: 's1', id: 'old', body: 'x'.repeat(4096), meta: { readOnly: false } }, + { storageSessionId: 's1', id: 'new', body: 'y'.repeat(4096), meta: { readOnly: false } }, + ]); + await storeCachedInputs(Readable.from(batch)); + const oldKey = inputCacheKey('s1', 'old'); + /* Backdate the older entry so LRU ordering is deterministic. */ + const past = new Date(Date.now() - 60_000); + await fsp.utimes(path.join(SESSION_INPUT_CACHE_DIR, oldKey), past, past); + await fsp.utimes(path.join(SESSION_INPUT_CACHE_DIR, `${oldKey}.json`), past, past); + + await pruneInputCache(5000); + + expect(await hasCachedInput('s1', 'old')).toBe(false); + expect(await hasCachedInput('s1', 'new')).toBe(true); + /* The sidecar must go with its object, or metadata leaks forever. */ + expect( + await fsp.lstat(path.join(SESSION_INPUT_CACHE_DIR, `${oldKey}.json`)).catch(() => null), + ).toBeNull(); + }); + + test('makes staging room from cold cache entries before committing a missing batch', async () => { + const initial = await makeBatch([ + { storageSessionId: 's1', id: 'cold', body: 'c'.repeat(4096), meta: { readOnly: false } }, + { storageSessionId: 's1', id: 'working', body: 'w'.repeat(4096), meta: { readOnly: false } }, + ]); + await storeCachedInputs(Readable.from(initial)); + + const coldKey = inputCacheKey('s1', 'cold'); + const past = new Date(Date.now() - 60_000); + await fsp.utimes(path.join(SESSION_INPUT_CACHE_DIR, coldKey), past, past); + await fsp.utimes(path.join(SESSION_INPUT_CACHE_DIR, `${coldKey}.json`), past, past); + /* A probe hit explicitly refreshes LRU state even on relatime tmpfs. */ + expect(await hasCachedInput('s1', 'working')).toBe(true); + + const incoming = await makeBatch([ + { storageSessionId: 's1', id: 'missing', body: 'm'.repeat(4096), meta: { readOnly: false } }, + ]); + /* Only two objects fit this logical cache budget. storeCachedInputs itself + * must evict the cold pair while staging, before the route's post-store + * prune can run. */ + await storeCachedInputs(Readable.from(incoming), 8500); + + expect(await hasCachedInput('s1', 'cold')).toBe(false); + expect(await hasCachedInput('s1', 'working')).toBe(true); + expect(await hasCachedInput('s1', 'missing')).toBe(true); + }); + + test('rejects a batch whose declared expanded size does not match its members', async () => { + const batch = await makeBatch([ + { storageSessionId: 's1', id: 'mismatch', body: 'bytes', meta: { readOnly: false } }, + ]); + await expect( + storeCachedInputs(Readable.from(batch), 1024, 1), + ).rejects.toThrow('expanded-byte mismatch'); + expect(await hasCachedInput('s1', 'mismatch')).toBe(false); + }); +}); diff --git a/api/src/session-inputs.ts b/api/src/session-inputs.ts new file mode 100644 index 0000000..0f1b7f3 --- /dev/null +++ b/api/src/session-inputs.ts @@ -0,0 +1,657 @@ +import * as crypto from 'crypto'; +import * as fs from 'fs'; +import * as fsp from 'fs/promises'; +import * as path from 'path'; +import { Transform } from 'stream'; +import { createGunzip } from 'zlib'; +import { logger } from './logger'; +import { SANDBOX_WORKSPACE_ROOT } from './workspace-isolation'; + +/** + * Runner-local cache of by-reference input objects. + * + * Backends whose sandbox cannot reach the file server (the Lambda MicroVM's + * only egress is the public internet) have the control plane PUSH input bytes + * into this cache before an execute. Priming then resolves a ref from here + * instead of over HTTP — see `Job.fetchInputObject`. + * + * The cache is deliberately NOT part of the session workspace: + * - entries are keyed by a runner-computed digest of (storage session, id), + * so no caller-supplied path component ever reaches the filesystem; + * - sandbox code cannot see or modify it (0700, outside /mnt/data); + * - it is never checkpointed, so a relaunched VM simply starts empty and the + * control plane re-pushes what the next execute needs. + * + * Because the only workspace writer stays the existing prime path, delivered + * files inherit its identity, ownership, read-only, symlink-safety, priming + * and modification-detection semantics for free. + */ + +/** + * Sibling of the workspace root, never inside it. Everything under + * SANDBOX_WORKSPACE_ROOT is a workspace as far as the stale-workspace reaper + * is concerned, so a cache placed there is deleted out from under the very + * execute it was pushed for — proven live, then reproduced with + * `reapStaleWorkspaces`. Being outside also keeps it off every nsjail mount, + * so sandbox code can neither read nor tamper with pending inputs. + */ +const DEFAULT_SESSION_INPUT_CACHE_DIR = `${SANDBOX_WORKSPACE_ROOT}-inputs`; +const UNSAFE_CACHE_ROOTS = new Set([ + '/', + '/bin', + '/boot', + '/dev', + '/etc', + '/home', + '/lib', + '/lib64', + '/mnt', + '/mnt/data', + '/opt', + '/proc', + '/root', + '/run', + '/sbin', + '/srv', + '/sys', + '/tmp', + '/usr', + '/var', + '/var/tmp', +]); + +/** + * The cache is root-owned and recursively pruned. Refuse broad paths and any + * path that contains, or is contained by, the workspace root so a typo cannot + * turn cache eviction into workspace deletion. + */ +export function resolveSessionInputCacheDir( + configured = process.env.SANDBOX_INPUT_CACHE_DIR ?? DEFAULT_SESSION_INPUT_CACHE_DIR, +): string { + if (!path.isAbsolute(configured)) { + throw new Error('SANDBOX_INPUT_CACHE_DIR must be an absolute path'); + } + const resolved = path.resolve(configured); + const workspace = path.resolve(SANDBOX_WORKSPACE_ROOT); + const cacheFromWorkspace = path.relative(workspace, resolved); + const workspaceFromCache = path.relative(resolved, workspace); + const isInsideWorkspace = + cacheFromWorkspace === '' || + (!cacheFromWorkspace.startsWith(`..${path.sep}`) && cacheFromWorkspace !== '..'); + const containsWorkspace = + workspaceFromCache === '' || + (!workspaceFromCache.startsWith(`..${path.sep}`) && workspaceFromCache !== '..'); + if (UNSAFE_CACHE_ROOTS.has(resolved) || isInsideWorkspace || containsWorkspace) { + throw new Error( + `SANDBOX_INPUT_CACHE_DIR must be a dedicated directory outside ${SANDBOX_WORKSPACE_ROOT}`, + ); + } + return resolved; +} + +export const SESSION_INPUT_CACHE_DIR = resolveSessionInputCacheDir(); +/** Cache entry names are exactly this shape — anything else is rejected. */ +const ENTRY_PATTERN = /^[0-9a-f]{64}(\.json)?$/; +const META_SUFFIX = '.json'; +const MAX_META_BYTES = 1024; +const TAR_BLOCK_BYTES = 512; +export const SESSION_INPUT_CACHE_MAX_OBJECTS = 256; +const TAR_MEMBER_LIMIT = SESSION_INPUT_CACHE_MAX_OBJECTS * 4 + 16; +/* Payload bytes have their own exact maxBytes accounting. Allow at most one + * header plus one padding block per legal member, plus terminators. Counting + * the whole decompressed stream against this bound closes gzip bombs made from + * unlimited zero blocks after an otherwise-valid tar terminator. */ +const TAR_STRUCTURAL_OVERHEAD_BYTES = (TAR_MEMBER_LIMIT * 2 + 4) * TAR_BLOCK_BYTES; + +type WritableFileHandle = Pick; + +/** FileHandle.write may legally complete with a short write. Consume the + * entire slice or fail instead of treating unwritten bytes as cached input. */ +export async function writeAllToHandle( + handle: WritableFileHandle, + data: Buffer, + recoverNoSpace?: () => Promise, +): Promise { + let offset = 0; + let recovered = false; + while (offset < data.length) { + let bytesWritten: number; + try { + ({ bytesWritten } = await handle.write( + data, + offset, + data.length - offset, + null, + )); + } catch (error) { + if (!recovered && recoverNoSpace && isNoSpaceError(error)) { + recovered = true; + await recoverNoSpace(); + continue; + } + throw error; + } + if (bytesWritten <= 0) { + throw new Error('Session input cache write made no progress'); + } + offset += bytesWritten; + } +} + +export interface CachedInputMeta { + /** Whether the object is infrastructure the sandbox must not modify. This is + * a property of the OBJECT, so it belongs here. A filename deliberately + * does not: the cache is keyed by object, and the same object can be + * requested at several destinations in one execute, so the requesting ref + * owns the name. Emitting one cached name as Content-Disposition made every + * ref resolve to the first ref's path — which then overwrote a file the + * sandbox had edited. */ + readOnly: boolean; +} + +export interface CachedInput { + path: string; + meta: CachedInputMeta; + /** Held open with O_NOFOLLOW so eviction can unlink the cache pathname + * without invalidating an execute that already accepted this entry. */ + handle: fsp.FileHandle; +} + +/** Opaque, collision-resistant key for a (storage session, object) pair. The + * digest keeps caller-controlled ids out of the filesystem entirely. */ +export function inputCacheKey(storageSessionId: string, id: string): string { + return crypto + .createHash('sha256') + .update(`${storageSessionId}\u0000${id}`, 'utf8') + .digest('hex'); +} + +function parseCachedInputMeta(raw: string): CachedInputMeta | null { + try { + const parsed = JSON.parse(raw) as unknown; + if ( + typeof parsed !== 'object' || + parsed === null || + typeof (parsed as { readOnly?: unknown }).readOnly !== 'boolean' + ) { + return null; + } + return { readOnly: (parsed as { readOnly: boolean }).readOnly }; + } catch { + return null; + } +} + +export async function hasCachedInput( + storageSessionId: string, + id: string, + cacheKey?: string, +): Promise { + const opened = await openCachedInput(storageSessionId, id, cacheKey); + if (!opened) return false; + await opened.handle.close(); + return true; +} + +export async function openCachedInput( + storageSessionId: string, + id: string, + cacheKey?: string, +): Promise { + const key = cacheKey ?? inputCacheKey(storageSessionId, id); + if (!/^[0-9a-f]{64}$/.test(key)) return null; + const entry = path.join(SESSION_INPUT_CACHE_DIR, key); + let handle: fsp.FileHandle | undefined; + try { + handle = await fsp.open(entry, fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW); + const dataStat = await handle.stat(); + if (!dataStat.isFile()) { + await handle.close(); + return null; + } + + let metaHandle: fsp.FileHandle | undefined; + let metaMtime: Date | undefined; + let raw: string | null = null; + try { + metaHandle = await fsp.open( + `${entry}${META_SUFFIX}`, + fs.constants.O_RDONLY | fs.constants.O_NOFOLLOW, + ); + const metaStat = await metaHandle.stat(); + metaMtime = metaStat.mtime; + if (metaStat.isFile() && metaStat.size <= MAX_META_BYTES) { + raw = await metaHandle.readFile('utf8'); + } + } catch { + raw = null; + } finally { + await metaHandle?.close().catch(() => {}); + } + const meta = raw === null ? null : parseCachedInputMeta(raw); + if (!meta) { + logger.warn({ key }, 'Ignoring session input with missing or invalid metadata'); + await handle.close(); + return null; + } + /* tmpfs commonly uses relatime, so merely opening a hit does not reliably + * refresh atime. Touch both immutable cache files explicitly: capacity + * recovery can then evict cold entries while preserving the working set + * the immediately preceding probe proved present. */ + const now = new Date(); + await Promise.all([ + fsp.utimes(entry, now, dataStat.mtime).catch(() => {}), + fsp.utimes(`${entry}${META_SUFFIX}`, now, metaMtime ?? now).catch(() => {}), + ]); + return { path: entry, meta, handle }; + } catch { + await handle?.close().catch(() => {}); + return null; + } +} + +/** + * Presents a cached entry as the `Response` the file server would have + * returned, so every downstream priming step (name resolution, read-only + * protection, streaming hash, atomic rename) runs byte-identically whether the + * bytes arrived by pull or by push. + */ +export function cachedInputResponse(entry: CachedInput): Response { + const headers = new Headers(); + /* No Content-Disposition: priming falls back to the ref's requested name, + * so each destination gets its own copy (see CachedInputMeta). */ + if (entry.meta.readOnly === true) headers.set('x-read-only', 'true'); + return new Response(entry.handle.createReadStream() as unknown as ReadableStream, { + status: 200, + headers, + }); +} + +function tarString(block: Buffer, offset: number, length: number): string { + const field = block.subarray(offset, offset + length); + const nul = field.indexOf(0); + return field.subarray(0, nul === -1 ? field.length : nul).toString('utf8'); +} + +function tarOctal(block: Buffer, offset: number, length: number): number { + const raw = tarString(block, offset, length).trim().replace(/\0/g, ''); + if (!/^[0-7]*$/.test(raw)) throw new Error('Invalid tar numeric field'); + return raw === '' ? 0 : Number.parseInt(raw, 8); +} + +function validateTarChecksum(block: Buffer): void { + const expected = tarOctal(block, 148, 8); + let actual = 0; + for (let i = 0; i < block.length; i += 1) { + actual += i >= 148 && i < 156 ? 0x20 : block[i]; + } + if (actual !== expected) throw new Error('Invalid session input tar checksum'); +} + +/** + * Extracts the deliberately tiny flat tar contract without invoking `tar`. + * This lets the runner enforce member count and expanded bytes before writes, + * reject links/PAX/path traversal, and cap compressed input while streaming. + */ +async function extractInputArchive( + body: NodeJS.ReadableStream, + staging: string, + maxBytes: number, + reserveStagingBytes: (bytes: number) => Promise, + recoverNoSpace: () => Promise, +): Promise { + let compressedBytes = 0; + const compressedLimit = maxBytes + Math.max(1024 * 1024, Math.ceil(maxBytes / 100)); + const decompressedLimit = maxBytes > Number.MAX_SAFE_INTEGER - TAR_STRUCTURAL_OVERHEAD_BYTES + ? Number.MAX_SAFE_INTEGER + : maxBytes + TAR_STRUCTURAL_OVERHEAD_BYTES; + const compressedGuard = new Transform({ + transform(chunk: Buffer, _encoding, callback) { + compressedBytes += chunk.length; + callback( + compressedBytes > compressedLimit + ? new Error(`Compressed session input batch exceeds ${compressedLimit} bytes`) + : null, + chunk, + ); + }, + }); + const gunzip = createGunzip(); + const forwardCompressedError = (error: Error): void => { + gunzip.destroy(error); + }; + const forwardBodyError = (error: unknown): void => { + compressedGuard.destroy(error instanceof Error ? error : new Error(String(error))); + }; + compressedGuard.on('error', forwardCompressedError); + body.once('error', forwardBodyError); + body.pipe(compressedGuard).pipe(gunzip); + + let buffered = Buffer.alloc(0); + let current: + | { handle?: fsp.FileHandle; remaining: number; padding: number; expectedSize?: number } + | undefined; + let expandedBytes = 0; + let files = 0; + let members = 0; + let decompressedBytes = 0; + let stagedBytes = 0; + let archiveEnded = false; + const seen = new Set(); + + try { + for await (const rawChunk of gunzip) { + decompressedBytes += (rawChunk as Buffer).length; + if (decompressedBytes > decompressedLimit) { + throw new Error( + `Decompressed session input batch exceeds ${decompressedLimit} bytes`, + ); + } + let chunk = rawChunk as Buffer; + while (chunk.length > 0) { + if (archiveEnded) { + if (chunk.some(byte => byte !== 0)) { + throw new Error('Unexpected data after session input tar terminator'); + } + break; + } + + if (current) { + if (current.remaining > 0) { + const length = Math.min(current.remaining, chunk.length); + if (current.handle) { + await writeAllToHandle( + current.handle, + chunk.subarray(0, length), + recoverNoSpace, + ); + } + current.remaining -= length; + chunk = chunk.subarray(length); + if (current.remaining > 0) continue; + } + if (current.padding > 0) { + const length = Math.min(current.padding, chunk.length); + current.padding -= length; + chunk = chunk.subarray(length); + if (current.padding > 0) continue; + } + if (current.handle) { + const stat = await current.handle.stat(); + if (stat.size !== current.expectedSize) { + throw new Error( + `Session input cache write size mismatch: expected ${current.expectedSize}, got ${stat.size}`, + ); + } + await current.handle.close(); + } + current = undefined; + continue; + } + + const needed = TAR_BLOCK_BYTES - buffered.length; + const length = Math.min(needed, chunk.length); + buffered = Buffer.concat([buffered, chunk.subarray(0, length)]); + chunk = chunk.subarray(length); + if (buffered.length < TAR_BLOCK_BYTES) continue; + + const header = buffered; + buffered = Buffer.alloc(0); + if (header.every(byte => byte === 0)) { + archiveEnded = true; + continue; + } + members += 1; + if (members > TAR_MEMBER_LIMIT) { + throw new Error('Session input tar contains too many members'); + } + validateTarChecksum(header); + const prefix = tarString(header, 345, 155); + const rawName = [prefix, tarString(header, 0, 100)].filter(Boolean).join('/'); + const name = rawName.replace(/^\.\//, ''); + const size = tarOctal(header, 124, 12); + const type = String.fromCharCode(header[156] || 0x30); + + /* BSD tar emits per-file PAX metadata even for short names. Ignore the + * metadata payload completely; unlike a general tar extractor, this + * parser never applies PAX path/link overrides. It still counts toward + * the expanded-byte ceiling. */ + if (type === 'x' || type === 'g') { + expandedBytes += size; + if (expandedBytes > maxBytes) { + throw new Error(`Session input batch exceeds the ${maxBytes}-byte cache limit`); + } + current = { + remaining: size, + padding: (TAR_BLOCK_BYTES - (size % TAR_BLOCK_BYTES)) % TAR_BLOCK_BYTES, + }; + continue; + } + if (type === '5' && (name === '' || name === '.')) { + if (size !== 0) throw new Error('Invalid tar directory member size'); + continue; + } + if (type !== '0') { + throw new Error(`Unsupported session input tar member type: ${type}`); + } + if (!ENTRY_PATTERN.test(name) || seen.has(name)) { + throw new Error(`Unexpected session input member: ${name}`); + } + seen.add(name); + files += 1; + if (files > SESSION_INPUT_CACHE_MAX_OBJECTS * 2) { + throw new Error( + `Session input batch exceeds the ${SESSION_INPUT_CACHE_MAX_OBJECTS}-object limit`, + ); + } + if (name.endsWith(META_SUFFIX) && size > MAX_META_BYTES) { + throw new Error(`Session input metadata exceeds ${MAX_META_BYTES} bytes`); + } + expandedBytes += size; + if (expandedBytes > maxBytes) { + throw new Error(`Session input batch exceeds the ${maxBytes}-byte cache limit`); + } + stagedBytes += size; + await reserveStagingBytes(stagedBytes); + let handle: fsp.FileHandle; + try { + handle = await fsp.open(path.join(staging, name), 'wx', 0o600); + } catch (error) { + if (!isNoSpaceError(error)) throw error; + await recoverNoSpace(); + handle = await fsp.open(path.join(staging, name), 'wx', 0o600); + } + current = { + handle, + remaining: size, + padding: (TAR_BLOCK_BYTES - (size % TAR_BLOCK_BYTES)) % TAR_BLOCK_BYTES, + expectedSize: size, + }; + } + } + if (!archiveEnded || current || buffered.length !== 0) { + throw new Error('Truncated session input tar archive'); + } + return stagedBytes; + } finally { + body.removeListener('error', forwardBodyError); + compressedGuard.removeListener('error', forwardCompressedError); + await current?.handle?.close().catch(() => {}); + if (!compressedGuard.destroyed) compressedGuard.destroy(); + if (!gunzip.destroyed) gunzip.destroy(); + } +} + +function isNoSpaceError(error: unknown): boolean { + return ( + typeof error === 'object' + && error !== null + && 'code' in error + && (error as { code?: unknown }).code === 'ENOSPC' + ); +} + +async function cleanupStaleInputStaging(): Promise { + const entries = await fsp.readdir(SESSION_INPUT_CACHE_DIR, { withFileTypes: true }).catch(() => []); + const staleBefore = Date.now() - 10 * 60_000; + await Promise.all(entries + .filter(entry => entry.isDirectory() && entry.name.startsWith('.staging-')) + .map(async entry => { + const target = path.join(SESSION_INPUT_CACHE_DIR, entry.name); + const stat = await fsp.lstat(target).catch(() => null); + if (stat && stat.mtimeMs < staleBefore) { + await fsp.rm(target, { recursive: true, force: true }).catch(() => {}); + } + })); +} + +/** + * Extracts a pushed batch into the cache. Members are digest-named by the + * control plane (`` for bytes, `.json` for metadata), so validation + * is a name-shape check rather than path arithmetic: nothing else can be + * written, and no member can escape the cache directory. + */ +let storeTail: Promise = Promise.resolve(); + +async function storeCachedInputsOnce( + body: NodeJS.ReadableStream, + maxBytes = Number.MAX_SAFE_INTEGER, + expectedBytes?: number, +): Promise { + await fsp.mkdir(SESSION_INPUT_CACHE_DIR, { recursive: true, mode: 0o700 }); + const cacheStat = await fsp.lstat(SESSION_INPUT_CACHE_DIR); + if (!cacheStat.isDirectory() || cacheStat.isSymbolicLink()) { + throw new Error('Session input cache root must be a real directory'); + } + await fsp.chmod(SESSION_INPUT_CACHE_DIR, 0o700); + await cleanupStaleInputStaging(); + const staging = await fsp.mkdtemp(path.join(SESSION_INPUT_CACHE_DIR, '.staging-')); + try { + /* Treat the configured cache ceiling as a budget shared by committed + * entries and the in-flight staging tree. As each tar header declares its + * file size, evict cold entries BEFORE writing those bytes; this prevents + * a full cache from requiring a second cache-sized allocation just to + * receive a valid missing batch. An ENOSPC caused by unrelated workspace + * growth gets one emergency all-cache prune and exact-operation retry. */ + const stagedBytes = await extractInputArchive( + body, + staging, + maxBytes, + bytes => pruneInputCache(Math.max(0, maxBytes - bytes)), + () => pruneInputCache(0), + ); + if (expectedBytes !== undefined && stagedBytes !== expectedBytes) { + throw new Error( + `Session input expanded-byte mismatch: expected ${expectedBytes}, got ${stagedBytes}`, + ); + } + + const entries = await fsp.readdir(staging, { withFileTypes: true }); + for (const entry of entries) { + if (!entry.isFile() || !ENTRY_PATTERN.test(entry.name)) { + throw new Error(`Unexpected session input member: ${entry.name}`); + } + } + + const names = new Set(entries.map(entry => entry.name)); + const keys = entries + .filter(entry => !entry.name.endsWith(META_SUFFIX)) + .map(entry => entry.name); + if (keys.length * 2 !== entries.length) { + throw new Error('Each session input requires exactly one data file and metadata sidecar'); + } + for (const key of keys) { + const metaName = `${key}${META_SUFFIX}`; + if (!names.has(metaName)) { + throw new Error(`Session input ${key} is missing metadata`); + } + const raw = await fsp.readFile(path.join(staging, metaName), 'utf8'); + if (!parseCachedInputMeta(raw)) { + throw new Error(`Session input ${key} has invalid metadata`); + } + } + + let stored = 0; + /* Commit sidecars before data. A new key remains a probe miss until both + * exist; replacing an immutable key can only expose its new validated + * metadata beside identical object bytes for the brief rename window. */ + const commitOrder = [ + ...entries.filter(entry => entry.name.endsWith(META_SUFFIX)), + ...entries.filter(entry => !entry.name.endsWith(META_SUFFIX)), + ]; + for (const entry of commitOrder) { + await fsp.rename( + path.join(staging, entry.name), + path.join(SESSION_INPUT_CACHE_DIR, entry.name), + ); + if (!entry.name.endsWith(META_SUFFIX)) stored += 1; + } + return stored; + } finally { + await fsp.rm(staging, { recursive: true, force: true }).catch(() => {}); + } +} + +export async function storeCachedInputs( + body: NodeJS.ReadableStream, + maxBytes = Number.MAX_SAFE_INTEGER, + expectedBytes?: number, +): Promise { + /* Concurrent pushes otherwise each budget only its own staging tree and can + * collectively recreate the same transient disk spike. Queue extraction; + * request streams naturally backpressure while waiting. */ + const previous = storeTail; + let release!: () => void; + storeTail = new Promise((resolve) => { + release = resolve; + }); + await previous; + try { + return await storeCachedInputsOnce(body, maxBytes, expectedBytes); + } finally { + release(); + } +} + +/** Drops least-recently-used entries until the cache fits `maxBytes`. Eviction + * is always safe: a miss simply re-pushes on the next probe. */ +export async function pruneInputCache(maxBytes: number): Promise { + const names = await fsp.readdir(SESSION_INPUT_CACHE_DIR).catch(() => [] as string[]); + const nameSet = new Set(names.filter(name => ENTRY_PATTERN.test(name))); + const pairs: Array<{ key: string; size: number; atime: number }> = []; + let total = 0; + for (const name of nameSet) { + if (name.endsWith(META_SUFFIX)) continue; + const metaName = `${name}${META_SUFFIX}`; + const dataStat = await fsp.lstat(path.join(SESSION_INPUT_CACHE_DIR, name)).catch(() => null); + const metaStat = nameSet.has(metaName) + ? await fsp.lstat(path.join(SESSION_INPUT_CACHE_DIR, metaName)).catch(() => null) + : null; + if (!dataStat?.isFile() || !metaStat?.isFile()) { + await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, name), { force: true }).catch(() => {}); + await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, metaName), { force: true }).catch(() => {}); + nameSet.delete(metaName); + continue; + } + const size = dataStat.size + metaStat.size; + pairs.push({ key: name, size, atime: Math.max(dataStat.atimeMs, metaStat.atimeMs) }); + total += size; + nameSet.delete(metaName); + } + /* Metadata with no data object is never usable and must not accumulate. */ + for (const orphan of nameSet) { + if (orphan.endsWith(META_SUFFIX)) { + await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, orphan), { force: true }).catch(() => {}); + } + } + if (total <= maxBytes) return; + pairs.sort((a, b) => a.atime - b.atime); + for (const pair of pairs) { + if (total <= maxBytes) break; + await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, pair.key), { force: true }).catch(() => {}); + await fsp + .rm(path.join(SESSION_INPUT_CACHE_DIR, `${pair.key}${META_SUFFIX}`), { force: true }) + .catch(() => {}); + total -= pair.size; + } +} diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts new file mode 100644 index 0000000..adb10f7 --- /dev/null +++ b/api/src/session-workspace.test.ts @@ -0,0 +1,169 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { config } from './config'; +import { + SessionWorkspace, + bindSessionWorkspace, + getBoundSessionWorkspace, + parseSessionBinding, + parseSessionBindingFromHeader, + resetSessionWorkspaceStateForTests, + unbindSessionWorkspace, +} from './session-workspace'; + +const savedEnabled = config.session_workspace_enabled; + +afterEach(async () => { + config.session_workspace_enabled = savedEnabled; + await unbindSessionWorkspace().catch(() => {}); + resetSessionWorkspaceStateForTests(); +}); + +describe('parseSessionBinding (gating)', () => { + test('returns undefined when the image-level flag is off, regardless of payload', () => { + config.session_workspace_enabled = false; + expect(parseSessionBinding(JSON.stringify({ runtime_session_id: 'rt_1', session_workspace: true }))).toBeUndefined(); + }); + + test('binds only when enabled AND the payload opts in with a runtime_session_id', () => { + config.session_workspace_enabled = true; + expect(parseSessionBinding(JSON.stringify({ runtime_session_id: 'rt_1', session_workspace: true }))) + .toEqual({ runtimeSessionId: 'rt_1' }); + }); + + test('rejects payloads missing the opt-in flag or the session id', () => { + config.session_workspace_enabled = true; + expect(parseSessionBinding(JSON.stringify({ runtimeSessionId: 'rt_1' }))).toBeUndefined(); + expect(parseSessionBinding(JSON.stringify({ session_workspace: true }))).toBeUndefined(); + expect(parseSessionBinding(JSON.stringify({ session_workspace: true, runtime_session_id: '' }))).toBeUndefined(); + }); +}); + +describe('parseSessionBindingFromHeader (per-request opt-in)', () => { + test('returns undefined when the image-level flag is off', () => { + config.session_workspace_enabled = false; + expect(parseSessionBindingFromHeader('rt_abc123')).toBeUndefined(); + }); + + test('binds a well-formed id when enabled (presence of the header is the opt-in)', () => { + config.session_workspace_enabled = true; + expect(parseSessionBindingFromHeader('rt_abc123')).toEqual({ runtimeSessionId: 'rt_abc123' }); + expect(parseSessionBindingFromHeader(' rt_abc123 ')).toEqual({ runtimeSessionId: 'rt_abc123' }); + }); + + test('rejects missing, empty, repeated, or malformed headers', () => { + config.session_workspace_enabled = true; + expect(parseSessionBindingFromHeader(undefined)).toBeUndefined(); + expect(() => parseSessionBindingFromHeader('')).toThrow('malformed'); + expect(() => parseSessionBindingFromHeader(['rt_a', 'rt_b'])).toThrow('exactly once'); + expect(() => parseSessionBindingFromHeader('rt bad space')).toThrow('malformed'); + expect(() => parseSessionBindingFromHeader('a'.repeat(129))).toThrow('malformed'); + }); + + test('tolerates absent and non-JSON payloads', () => { + config.session_workspace_enabled = true; + expect(parseSessionBinding(undefined)).toBeUndefined(); + expect(parseSessionBinding('')).toBeUndefined(); + expect(parseSessionBinding('not json')).toBeUndefined(); + }); +}); + +describe('bindSessionWorkspace lifecycle', () => { + test('binding is idempotent for the same runtime session and returns the same instance', () => { + const a = bindSessionWorkspace({ runtimeSessionId: 'rt_1' }); + const b = bindSessionWorkspace({ runtimeSessionId: 'rt_1' }); + expect(a).toBe(b); + expect(getBoundSessionWorkspace()).toBe(a); + }); + + /* One runner serves exactly one session for its lifetime: honoring a second + * id would race the previous session's async wipe against the new session's + * restore over the same directory. The bind must fail closed instead. */ + test('a different runtime session is rejected, never rebound', () => { + const a = bindSessionWorkspace({ runtimeSessionId: 'rt_1' }); + const b = bindSessionWorkspace({ runtimeSessionId: 'rt_2' }); + expect(b).toBeUndefined(); + expect(getBoundSessionWorkspace()).toBe(a); + expect(getBoundSessionWorkspace()?.runtimeSessionId).toBe('rt_1'); + }); + + test('unbind clears the bound session', async () => { + bindSessionWorkspace({ runtimeSessionId: 'rt_1' }); + await unbindSessionWorkspace(); + expect(getBoundSessionWorkspace()).toBeUndefined(); + }); +}); + +describe('SessionWorkspace state', () => { + test('surfaced and primed tracking, cleared on reset', async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), 'sw-state-')); + const savedPerJob = config.per_job_uids; + config.per_job_uids = false; + try { + const ws = new SessionWorkspace({ runtimeSessionId: 'rt_1' }); + + expect(ws.isSurfaced('out.csv', '10:100')).toBe(false); + ws.markSurfaced('out.csv', '10:100'); + expect(ws.isSurfaced('out.csv', '10:100')).toBe(true); + expect(ws.isSurfaced('out.csv', '11:200')).toBe(false); + + expect(ws.primedInputId('in.csv')).toBeUndefined(); + expect(ws.isPrimedInput('in.csv')).toBe(false); + ws.markPrimed('in.csv', 'file_abc'); + expect(ws.primedInputId('in.csv')).toBe('file_abc'); + + /* read-only primes report as not-primed so the caller re-downloads them + * (a reused on-disk copy could have been tampered via the writable dir). */ + ws.markPrimed('skill.py', 'file_ro', true); + expect(ws.primedInputId('skill.py')).toBeUndefined(); + /* ...but both still count as primed inputs, so a later turn that omits + * them doesn't re-surface them as generated outputs. */ + expect(ws.isPrimedInput('in.csv')).toBe(true); + expect(ws.isPrimedInput('skill.py')).toBe(true); + expect(ws.isPrimedInput('never-primed.csv')).toBe(false); + /* read-only primes are always suppressed (modifications dropped by + * contract); writable ones are only suppressed while unchanged. */ + expect(ws.isPrimedReadOnly('skill.py')).toBe(true); + expect(ws.isPrimedReadOnly('in.csv')).toBe(false); + + await ws.reset(); + expect(ws.isSurfaced('out.csv', '10:100')).toBe(false); + expect(ws.primedInputId('in.csv')).toBeUndefined(); + } finally { + config.per_job_uids = savedPerJob; + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + test('snapshotMeta/loadMeta round-trips priming + output-diff state into a fresh workspace', () => { + const source = new SessionWorkspace({ runtimeSessionId: 'rt_1' }); + source.markSurfaced('out.csv', '10:100'); + source.markPrimed('in.csv', 'file_abc', false, 'ORIGHASH'); + source.markPrimed('skill.py', 'file_ro', true); + + /* A relaunched VM starts with an empty workspace and loads the checkpoint's + * sidecar — without this it would re-download every input, overwriting a + * restored in-place-modified file with the original. */ + const relaunched = new SessionWorkspace({ runtimeSessionId: 'rt_1' }); + relaunched.loadMeta(source.snapshotMeta()); + + expect(relaunched.primedInputId('in.csv')).toBe('file_abc'); + expect(relaunched.isSurfaced('out.csv', '10:100')).toBe(true); + /* the original upload hash survives so reuse baselines against it, not a + * re-hash of a possibly-mutated on-disk copy */ + expect(relaunched.primedHash('in.csv')).toBe('ORIGHASH'); + /* read-only flag survives the round-trip, so it still re-downloads */ + expect(relaunched.primedInputId('skill.py')).toBeUndefined(); + }); + + test('a partial prime stays fail-closed until a successful restore loads metadata', () => { + const workspace = new SessionWorkspace({ runtimeSessionId: 'rt_dirty' }); + workspace.markDirty('partial input prime'); + expect(workspace.dirtyReason).toBe('partial input prime'); + + workspace.loadMeta({ primed: [], surfaced: [] }); + expect(workspace.dirtyReason).toBeUndefined(); + }); +}); diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts new file mode 100644 index 0000000..e26fa05 --- /dev/null +++ b/api/src/session-workspace.ts @@ -0,0 +1,308 @@ +import { config } from './config'; +import { logger } from './logger'; +import { + ensureSessionWorkspace, + resetSessionWorkspace, + sandboxJobUidPool, + type SandboxJobIdentity, + type SandboxWorkspaceLease, +} from './workspace-isolation'; + +/** + * Persistent, stateful session workspace for the Lambda MicroVM backend. + * + * A session-bound VM runs exactly one runtime session, and its executions + * serialize on the control-plane lock — so the runner keeps a single + * long-lived workspace and a single pinned UID, reused by every `/execute`. + * This is what turns the semi-stateless runner stateful: files, installed + * packages, and chDB dirs under `/mnt/data` survive between calls instead of + * being wiped per job. + * + * Gated by two independent locks (both required): the image-level + * `SANDBOX_SESSION_WORKSPACE_ENABLED` (true only in the Lambda MicroVM runner + * target) and a per-request opt-in. The control plane opts a VM into session + * mode by stamping the derived runtime session id on every `/execute` via the + * `X-Runtime-Session-Id` header (see `parseSessionBindingFromHeader`). When + * neither lock is active, `getBoundSessionWorkspace()` returns undefined and + * the runner falls back to the untouched fresh-per-job path. + * + * The header, not a `/run` lifecycle hook, is the delivery mechanism: Lambda's + * image build hooks require the snapshot-compatible Lambda base container image + * to route, and enabling any runtime hook forces the `/ready` build hook, which + * never reaches a stock container's listener. Per-request signaling keeps image + * builds hookless (reliable) and needs no snapshot handshake. + */ + +/** Wire contract with the Lambda backend (`service/src/sandbox-backend`). */ +export const RUNTIME_SESSION_ID_HEADER = 'x-runtime-session-id'; + +/** Legacy checkpoint-sidecar path used before runner control metadata moved to + * its own top-level archive member. Restore still accepts this format for + * rolling compatibility; new checkpoints never create or alter this path. */ +export const SESSION_META_FILE = '.codeapi-session-meta.json'; + +/** Marker embedded in runner-owned checkpoint control metadata so restore can + * distinguish it from a user file that merely shares the legacy reserved name + * and happens to contain primed/surfaced arrays. */ +/* Bump whenever the meaning of a persisted field changes. v1 stored + * per-execution masked object handles; v2 stores the stable input-cache digest. + * Treating a v1 id as a v2 digest would re-prime over edited session inputs. */ +export const SESSION_META_MARKER = 'codeapi.session-meta.v2'; + +export interface SessionMetaSnapshot { + marker?: string; + primed: Array<[string, { id: string; readOnly: boolean; hash?: string }]>; + surfaced: Array<[string, string]>; +} + +const RUNTIME_SESSION_ID_PATTERN = /^[A-Za-z0-9._:-]{1,128}$/; + +export interface SessionBinding { + runtimeSessionId: string; +} + +export class SessionWorkspaceBindingError extends Error { + readonly status = 400; + + constructor(message: string) { + super(message); + this.name = 'SessionWorkspaceBindingError'; + } +} + +/** Shape of the `/run` runHookPayload the control plane delivers per VM. */ +interface RunHookSessionPayload { + runtime_session_id?: unknown; + session_workspace?: unknown; +} + +export function parseSessionBinding(runHookPayload: string | undefined): SessionBinding | undefined { + if (!config.session_workspace_enabled) return undefined; + if (runHookPayload == null || runHookPayload.length === 0) return undefined; + let parsed: RunHookSessionPayload; + try { + parsed = JSON.parse(runHookPayload) as RunHookSessionPayload; + } catch { + logger.warn('Ignoring non-JSON /run runHookPayload for session binding'); + return undefined; + } + if (parsed.session_workspace !== true) return undefined; + if (typeof parsed.runtime_session_id !== 'string' || parsed.runtime_session_id.length === 0) { + logger.warn('Session workspace requested without a runtime_session_id — ignoring'); + return undefined; + } + return { runtimeSessionId: parsed.runtime_session_id }; +} + +/** Per-request session opt-in from the `X-Runtime-Session-Id` header. Presence + * of a well-formed id is the opt-in; the header is only honored on the Lambda + * MicroVM runner target (`session_workspace_enabled`). Header values arrive as + * `string | string[]` from Node — a repeated header is malformed, so reject. */ +export function parseSessionBindingFromHeader( + headerValue: string | string[] | undefined, +): SessionBinding | undefined { + if (!config.session_workspace_enabled) return undefined; + if (headerValue === undefined) return undefined; + if (typeof headerValue !== 'string') { + throw new SessionWorkspaceBindingError('X-Runtime-Session-Id must appear exactly once'); + } + const runtimeSessionId = headerValue.trim(); + if (!RUNTIME_SESSION_ID_PATTERN.test(runtimeSessionId)) { + throw new SessionWorkspaceBindingError('X-Runtime-Session-Id is malformed'); + } + return { runtimeSessionId }; +} + +export class SessionWorkspace { + readonly runtimeSessionId: string; + private lease: SandboxWorkspaceLease | undefined; + private identity: SandboxJobIdentity | undefined; + /** relPath -> signature of every file already surfaced to the client, so a + * later job re-scanning the persistent workspace does not re-upload + * unchanged prior outputs (output diffing). */ + private readonly surfaced = new Map(); + /** relPath -> {id, readOnly, hash} already primed onto disk, so an unchanged + * input delivered again is not re-downloaded (priming dedup). `readOnly` + * inputs are never reused (a sandbox can unlink+replace a 0444 file via the + * writable parent dir), so they re-download to restore pristine content + + * protection. `hash` is the ORIGINAL upload hash, kept as the modification + * baseline on reuse so a writable input mutated by a prior turn is reported + * as modified-from-original rather than re-hashed as if it were pristine. */ + private readonly primed = new Map(); + /** A failed multi-file prime may have committed only part of the request. + * Keep the runner fail-closed until the control plane recycles it and restores + * the last committed checkpoint. This is intentionally in-memory only: a + * dirty workspace must never be checkpointed as a new recovery point. */ + private dirty: string | undefined; + + constructor(binding: SessionBinding) { + this.runtimeSessionId = binding.runtimeSessionId; + } + + /** Acquires (once) the pinned UID + persistent dir, reused every job. */ + async acquire(): Promise { + if (!this.identity) { + const identity = sandboxJobUidPool.acquire(); + if (!identity) { + throw new Error('No sandbox UID slot available for session workspace'); + } + this.identity = identity; + } + this.lease = await ensureSessionWorkspace(this.identity); + return this.lease; + } + + /** The pinned UID/GID for this session, so a restored checkpoint's files + * can be chowned to the owner the sandbox jobs run as. Ensures the + * workspace/identity exist first. */ + async ownership(): Promise<{ dir: string; uid: number; gid: number }> { + const lease = await this.acquire(); + return { dir: lease.dir, uid: lease.identity.uid, gid: lease.identity.gid }; + } + + isSurfaced(relPath: string, hash: string): boolean { + return this.surfaced.get(relPath) === hash; + } + + markSurfaced(relPath: string, hash: string): void { + this.surfaced.set(relPath, hash); + } + + forget(relPath: string): void { + this.surfaced.delete(relPath); + } + + /** The primed storage id for `relPath`, or undefined. `readOnly` primes are + * reported as not-primed so the caller always re-downloads them. */ + primedInputId(relPath: string): string | undefined { + const entry = this.primed.get(relPath); + if (!entry || entry.readOnly) return undefined; + return entry.id; + } + + /** Whether `relPath` was primed as an input on any earlier turn (regardless + * of read-only). Such a file persists in the workspace, so a later turn that + * doesn't re-send it must not mistake it for a newly generated output. */ + isPrimedInput(relPath: string): boolean { + return this.primed.has(relPath); + } + + /** Whether the primed input at `relPath` is read-only — its sandboxed-code + * modifications are dropped by contract and never surfaced as outputs. */ + isPrimedReadOnly(relPath: string): boolean { + return this.primed.get(relPath)?.readOnly === true; + } + + markPrimed(relPath: string, storageFileId: string, readOnly = false, hash?: string): void { + this.primed.set(relPath, { id: storageFileId, readOnly, hash }); + } + + markDirty(reason: string): void { + this.dirty = reason; + logger.error( + { runtimeSessionId: this.runtimeSessionId, reason }, + 'Session workspace marked dirty', + ); + } + + get dirtyReason(): string | undefined { + return this.dirty; + } + + /** The original upload hash recorded when `relPath` was first primed, or + * undefined. Used as the modification baseline on reuse. */ + primedHash(relPath: string): string | undefined { + return this.primed.get(relPath)?.hash; + } + + /** Serializes the priming + output-diff state into checkpoint control + * metadata so a relaunched VM restores the same in-memory state. */ + snapshotMeta(): SessionMetaSnapshot { + return { + primed: [...this.primed.entries()], + surfaced: [...this.surfaced.entries()], + }; + } + + /** Rebuilds priming + output-diff state from restored checkpoint metadata. + * Without it a relaunched VM would re-download every input ref, overwriting a + * restored in-place-modified input with its original, and re-upload every + * restored file as a new output. */ + loadMeta(snapshot: SessionMetaSnapshot): void { + /* REPLACE, never merge: the restored workspace is the authoritative + * content, so entries from a prior (failed or superseded) restore that the + * snapshot omits must not linger — a stale primed entry would suppress a + * real file from the output scan. */ + this.primed.clear(); + this.surfaced.clear(); + this.dirty = undefined; + for (const [relPath, entry] of snapshot.primed) this.primed.set(relPath, entry); + for (const [relPath, hash] of snapshot.surfaced) this.surfaced.set(relPath, hash); + } + + /** Full teardown: wipe the dir, release the pinned UID, clear diff state. + * Fail closed on a failed wipe: the directory was quarantined with this + * session's data still inside, so the pinned UID is deliberately NOT + * released — recycling it could let a later session reactivate the + * quarantined contents under a matching identity. Leaking one UID slot on + * a VM that is being recycled anyway is the cheaper failure. */ + async reset(): Promise { + const wiped = await resetSessionWorkspace(); + this.surfaced.clear(); + this.primed.clear(); + this.dirty = undefined; + this.lease = undefined; + if (!wiped) { + logger.error( + { runtimeSessionId: this.runtimeSessionId }, + 'Session workspace wipe failed; retaining pinned UID for the quarantined directory', + ); + return; + } + if (this.identity) { + sandboxJobUidPool.release(this.identity); + this.identity = undefined; + } + } +} + +let boundSession: SessionWorkspace | undefined; + +/** Binding the same session twice is a no-op. A DIFFERENT runtime session id + * is rejected outright: in the MicroVM topology one runner serves exactly one + * session for its whole lifetime, so a second id is a control-plane bug or a + * forged header — and honoring it would race an async wipe of the previous + * session against the new session's restore over the same directory (the + * reset could delete the new session's files, or checkpoint one session's + * data under the other's identity). Fail closed; the control plane recycles + * the VM on the resulting 409. */ +export function bindSessionWorkspace(binding: SessionBinding | undefined): SessionWorkspace | undefined { + if (!binding) return boundSession; + if (boundSession && boundSession.runtimeSessionId === binding.runtimeSessionId) { + return boundSession; + } + if (boundSession) { + logger.error( + { bound: boundSession.runtimeSessionId, requested: binding.runtimeSessionId }, + 'Refusing to rebind runner to a different runtime session', + ); + return undefined; + } + boundSession = new SessionWorkspace(binding); + return boundSession; +} + +export function getBoundSessionWorkspace(): SessionWorkspace | undefined { + return boundSession; +} + +/** Called by `/terminate` (and session reset) to release the workspace. */ +export async function unbindSessionWorkspace(): Promise { + const current = boundSession; + boundSession = undefined; + if (current) await current.reset(); +} + +export function resetSessionWorkspaceStateForTests(): void { + boundSession = undefined; +} diff --git a/api/src/tool-call-socket-process.test.ts b/api/src/tool-call-socket-process.test.ts new file mode 100644 index 0000000..3efe2a4 --- /dev/null +++ b/api/src/tool-call-socket-process.test.ts @@ -0,0 +1,119 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { config } from './config'; +import { + ensureToolCallSocketProxyReady, + stopToolCallSocketProxy, +} from './tool-call-socket-process'; + +const originalPort = config.allowed_local_network_port; +const originalEnv = { + target: process.env.SANDBOX_FORWARD_TARGET, + socket: process.env.TCS_SOCKET, + bundle: process.env.TCS_PROXY_BUNDLE, + node: process.env.TCS_NODE_BINARY, + marker: process.env.TCS_TEST_MARKER, +}; +let tempDir: string | undefined; + +function restoreEnv(name: string, value: string | undefined): void { + if (value === undefined) delete process.env[name]; + else process.env[name] = value; +} + +afterEach(async () => { + await stopToolCallSocketProxy(); + config.allowed_local_network_port = originalPort; + restoreEnv('SANDBOX_FORWARD_TARGET', originalEnv.target); + restoreEnv('TCS_SOCKET', originalEnv.socket); + restoreEnv('TCS_PROXY_BUNDLE', originalEnv.bundle); + restoreEnv('TCS_NODE_BINARY', originalEnv.node); + restoreEnv('TCS_TEST_MARKER', originalEnv.marker); + if (tempDir) await fsp.rm(tempDir, { recursive: true, force: true }); + tempDir = undefined; +}); + +describe('lazy tool-call socket process', () => { + test('starts once on demand, awaits readiness, and stops cleanly', async () => { + tempDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'lazy-tcs-')); + const socketPath = path.join(tempDir, 'proxy.sock'); + const markerPath = path.join(tempDir, 'launches.txt'); + const bundlePath = path.join(tempDir, 'fake-proxy.cjs'); + await fsp.writeFile(bundlePath, ` + const fs = require('fs'); + const net = require('net'); + try { fs.unlinkSync(process.env.TCS_SOCKET); } catch {} + fs.appendFileSync(process.env.TCS_TEST_MARKER, 'launch\\n'); + const server = net.createServer(socket => socket.end()); + server.listen(process.env.TCS_SOCKET); + const stop = () => server.close(() => process.exit(0)); + process.on('SIGTERM', stop); + process.on('SIGINT', stop); + `); + + config.allowed_local_network_port = 443; + process.env.SANDBOX_FORWARD_TARGET = 'https://gateway.example'; + process.env.TCS_SOCKET = socketPath; + process.env.TCS_PROXY_BUNDLE = bundlePath; + process.env.TCS_TEST_MARKER = markerPath; + + await Promise.all([ + ensureToolCallSocketProxyReady(), + ensureToolCallSocketProxyReady(), + ensureToolCallSocketProxyReady(), + ]); + + expect((await fsp.lstat(socketPath)).isSocket()).toBe(true); + expect((await fsp.readFile(markerPath, 'utf8')).trim().split('\n')).toHaveLength(1); + + /* A live child can lose its socket pathname (for example after an external + * cleanup). A previously-resolved launch promise must not mask that loss. */ + await fsp.unlink(socketPath); + await Promise.all([ + ensureToolCallSocketProxyReady(), + ensureToolCallSocketProxyReady(), + ensureToolCallSocketProxyReady(), + ]); + expect((await fsp.lstat(socketPath)).isSocket()).toBe(true); + expect((await fsp.readFile(markerPath, 'utf8')).trim().split('\n')).toHaveLength(2); + + await stopToolCallSocketProxy(); + expect(await fsp.lstat(socketPath).catch(() => null)).toBeNull(); + }); + + test('fails closed when forwarding is not configured', async () => { + config.allowed_local_network_port = 0; + process.env.SANDBOX_FORWARD_TARGET = ''; + await expect(ensureToolCallSocketProxyReady()).rejects.toThrow('not configured'); + }); + + test('reports a spawn failure without crashing and permits a later retry', async () => { + tempDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'lazy-tcs-spawn-')); + const socketPath = path.join(tempDir, 'proxy.sock'); + const bundlePath = path.join(tempDir, 'fake-proxy.cjs'); + await fsp.writeFile(bundlePath, ` + const fs = require('fs'); + const net = require('net'); + try { fs.unlinkSync(process.env.TCS_SOCKET); } catch {} + const server = net.createServer(socket => socket.end()); + server.listen(process.env.TCS_SOCKET); + process.on('SIGTERM', () => server.close(() => process.exit(0))); + `); + + config.allowed_local_network_port = 443; + process.env.SANDBOX_FORWARD_TARGET = 'https://gateway.example'; + process.env.TCS_SOCKET = socketPath; + process.env.TCS_PROXY_BUNDLE = bundlePath; + process.env.TCS_NODE_BINARY = path.join(tempDir, 'missing-node'); + + await expect(ensureToolCallSocketProxyReady()).rejects.toThrow( + 'tool-call socket proxy failed to spawn', + ); + + process.env.TCS_NODE_BINARY = process.execPath; + await ensureToolCallSocketProxyReady(); + expect((await fsp.lstat(socketPath)).isSocket()).toBe(true); + }); +}); diff --git a/api/src/tool-call-socket-process.ts b/api/src/tool-call-socket-process.ts new file mode 100644 index 0000000..4694c07 --- /dev/null +++ b/api/src/tool-call-socket-process.ts @@ -0,0 +1,191 @@ +import { spawn, type ChildProcess } from 'child_process'; +import * as net from 'net'; +import * as path from 'path'; +import { logger } from './logger'; +import { config } from './config'; + +const START_TIMEOUT_MS = 10_000; +const STOP_TIMEOUT_MS = 3_000; + +let child: ChildProcess | undefined; +let starting: Promise | undefined; + +function proxyBundlePath(): string { + return process.env.TCS_PROXY_BUNDLE + ?? path.resolve(__dirname, '..', '.build', 'tool-call-socket-proxy.cjs'); +} + +function proxySocketPath(): string { + return process.env.TCS_SOCKET || '/tmp/tcs.sock'; +} + +function proxyOwner(): { uid: string; gid: string } { + const runningAsRoot = typeof process.getuid === 'function' && process.getuid() === 0; + return { + uid: String(runningAsRoot ? 65534 : (process.getuid?.() ?? 65534)), + gid: String(runningAsRoot ? 65534 : (process.getgid?.() ?? 65534)), + }; +} + +function socketAcceptsConnections(socketPath: string): Promise { + return new Promise((resolve) => { + const socket = net.createConnection(socketPath); + let settled = false; + const finish = (ready: boolean): void => { + if (settled) return; + settled = true; + socket.destroy(); + resolve(ready); + }; + socket.once('connect', () => finish(true)); + socket.once('error', () => finish(false)); + socket.setTimeout(100, () => finish(false)); + }); +} + +async function waitUntilReady( + started: ChildProcess, + socketPath: string, + timeoutMs: number, +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (started.exitCode !== null || started.signalCode !== null) { + throw new Error( + `tool-call socket proxy exited before readiness` + + ` (code=${String(started.exitCode)}, signal=${String(started.signalCode)})`, + ); + } + if (await socketAcceptsConnections(socketPath)) return; + await new Promise(resolve => setTimeout(resolve, 25)); + } + throw new Error(`tool-call socket proxy did not become ready within ${timeoutMs}ms`); +} + +async function launch(): Promise { + const rawTarget = process.env.SANDBOX_FORWARD_TARGET?.trim() ?? ''; + if (config.allowed_local_network_port <= 0 || !rawTarget) { + throw new Error('tool-call socket proxy is not configured'); + } + + const socketPath = proxySocketPath(); + const owner = proxyOwner(); + const started = spawn(process.env.TCS_NODE_BINARY || 'node', [proxyBundlePath()], { + env: { + ...process.env, + TCS_SOCKET: socketPath, + TCS_SOCKET_UID: owner.uid, + TCS_SOCKET_GID: owner.gid, + SANDBOX_FORWARD_TARGET: rawTarget, + }, + stdio: ['ignore', 'pipe', 'pipe'], + }); + child = started; + let rejectSpawnFailure: (error: Error) => void; + const spawnFailure = new Promise((_resolve, reject) => { + rejectSpawnFailure = reject; + }); + /* A failed spawn emits `error`; it does not reliably set exitCode or + * signalCode. Without an immediate listener Node treats ENOENT/EACCES as an + * uncaught event and crashes the runner instead of failing this execution. */ + started.on('error', (error) => { + if (child === started) child = undefined; + const launchError = new Error( + `tool-call socket proxy failed to spawn: ${error.message}`, + { cause: error }, + ); + logger.error({ error }, 'tool-call socket proxy spawn failed'); + rejectSpawnFailure(launchError); + }); + started.stdout?.on('data', (chunk: Buffer) => { + logger.debug({ proxy: chunk.toString().trim() }, 'tool-call socket proxy'); + }); + started.stderr?.on('data', (chunk: Buffer) => { + logger.warn({ proxy: chunk.toString().trim() }, 'tool-call socket proxy stderr'); + }); + started.once('exit', (code, signal) => { + if (child === started) { + child = undefined; + } + if (code !== 0 && signal !== 'SIGTERM') { + logger.error({ code, signal }, 'tool-call socket proxy exited unexpectedly'); + } + }); + + try { + await Promise.race([ + waitUntilReady(started, socketPath, START_TIMEOUT_MS), + spawnFailure, + ]); + logger.info( + { socketPath, target: rawTarget }, + 'Tool-call socket proxy started after MicroVM restore', + ); + } catch (error) { + started.kill('SIGKILL'); + if (child === started) child = undefined; + throw error; + } +} + +/** + * Starts Node only when an authenticated execute actually receives the + * tool-call socket capability. Lambda MicroVM image creation snapshots the + * already-running container; starting Node in entrypoint would therefore clone + * its embedded OpenSSL process state into every VM. Lazy initialization keeps + * Node entirely on the post-restore side of that boundary. + */ +async function ensureToolCallSocketProxyReadyOnce(): Promise { + const socketPath = proxySocketPath(); + const running = child; + if (running && running.exitCode === null && running.signalCode === null) { + if (await socketAcceptsConnections(socketPath)) return; + running.kill('SIGKILL'); + if (child === running) child = undefined; + } + await launch(); +} + +export async function ensureToolCallSocketProxyReady(): Promise { + /* Serialize the whole probe/relaunch transaction, not only spawn readiness. + * Otherwise concurrent recovery callers can both await a failed socket probe, + * then each kill/relaunch through the shared `child` slot. Publishing this + * promise before the first probe await makes every follower join one attempt. */ + if (starting) { + await starting; + return; + } + const attempt = ensureToolCallSocketProxyReadyOnce(); + starting = attempt; + try { + await attempt; + } finally { + /* A successful launch must not leave a resolved promise cached forever: + * if the socket later disappears while the process is still alive, the + * next caller has to kill/relaunch rather than return false readiness. */ + if (starting === attempt) starting = undefined; + } +} + +export async function stopToolCallSocketProxy(): Promise { + const running = child; + child = undefined; + starting = undefined; + if (!running || running.exitCode !== null || running.signalCode !== null) return; + + await new Promise((resolve) => { + let settled = false; + const finish = (): void => { + if (settled) return; + settled = true; + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + running.kill('SIGKILL'); + finish(); + }, STOP_TIMEOUT_MS); + running.once('exit', finish); + running.kill('SIGTERM'); + }); +} diff --git a/api/src/walker.test.ts b/api/src/walker.test.ts index 74feea9..0dc47bd 100644 --- a/api/src/walker.test.ts +++ b/api/src/walker.test.ts @@ -8,6 +8,7 @@ import { Job, type TFile } from './job'; import type { Runtime } from './runtime'; import { config } from './config'; import { DIRKEEP } from './validation'; +import { SessionWorkspace } from './session-workspace'; /** * Integration tests for Job.walkDir and its extracted helpers. These tests @@ -25,6 +26,7 @@ interface WalkerInternals { generatedFiles: Array<{ id: string; name: string; path: string }>; sessionFiles: Array<{ id: string; name: string; storage_session_id: string; modified_from?: { id: string; storage_session_id: string }; inherited?: true; entity_id?: string }>; inheritedRefs: Array<{ id: string; name: string; storage_session_id: string; inherited?: true; entity_id?: string }>; + pendingSurfaced: Map; inputFileHashes: Map; files: TFile[]; walkDir: (dir: string, depth: number, inputByName: Map) => Promise<'collected' | 'empty' | 'skipped'>; @@ -61,6 +63,7 @@ function makeJob(opts: { * `storage_session_id` carried on each file ref. */ session_id?: string; maxFileSize?: number; + session?: SessionWorkspace; } = {}): Job { return new Job({ session_id: opts.session_id ?? 'test-session', @@ -71,6 +74,7 @@ function makeJob(opts: { timeouts: { compile: 5000, run: 5000 }, cpu_times: { compile: 5000, run: 5000 }, memory_limits: { compile: 100_000_000, run: 100_000_000 }, + session: opts.session, }); } @@ -372,13 +376,12 @@ describe('walkDir / regular file handling', () => { expect(internals.inheritedRefs[0]).not.toHaveProperty('entity_id'); }); - it('treats a downloaded input as generated when no hash baseline exists (e.g. download failed)', async () => { + it('treats a by-ref-shaped file without a hash baseline as generated', async () => { const name = 'inherited.py'; const full = path.join(tmpDir, name); - /* Simulate the failure scenario: downloadAndWriteFile returned null - * (never populated inputFileHashes), but the sandboxed run still - * produced bytes at this path. Echoing the stale inherited ref would - * lie to the client about content, so this must upload as generated. */ + /* Defense in depth for callers that construct a Job directly without + * priming it. Normal execution now fails before running when a required + * download cannot be completed. */ await fsp.writeFile(full, 'print("new bytes produced by the run")'); const downloaded: TFile = { @@ -808,6 +811,32 @@ describe('walkDir / dirent classification', () => { }); describe('createDirkeepMarker / symlink safety', () => { + it('tracks a successfully uploaded synthetic marker across session turns', async () => { + await fsp.mkdir(path.join(tmpDir, 'empty')); + const keepRel = path.join('empty', DIRKEEP); + const session = new SessionWorkspace({ runtimeSessionId: 'rt_dirkeep' }); + + const first = asInternals(makeJob({ session })); + first.submissionDir = tmpDir; + await first.walkDir(tmpDir, 0, new Map()); + + const marker = first.generatedFiles.find(file => file.name === keepRel); + expect(marker).toBeDefined(); + const pending = first.pendingSurfaced.get(marker!.id); + expect(pending).toEqual({ name: keepRel, signature: sha256('') }); + + /* uploadGeneratedFiles commits this exact pending entry only after the + * object PUT succeeds. Model that successful commit, then scan the same + * persistent workspace on the next turn. */ + session.markSurfaced(pending!.name, pending!.signature); + const second = asInternals(makeJob({ session })); + second.submissionDir = tmpDir; + await second.walkDir(tmpDir, 0, new Map()); + + expect(second.generatedFiles.some(file => file.name === keepRel)).toBe(false); + expect(second.sessionFiles.some(file => file.name === keepRel)).toBe(false); + }); + it('refuses to write a .dirkeep marker through an existing symlink', async () => { /* Simulate a malicious sandbox that plants a .dirkeep symlink pointing * to a sensitive file outside the empty directory. The marker write diff --git a/api/src/warmup.test.ts b/api/src/warmup.test.ts new file mode 100644 index 0000000..a3d7b9f --- /dev/null +++ b/api/src/warmup.test.ts @@ -0,0 +1,35 @@ +import { describe, expect, test } from 'bun:test'; +import * as fsp from 'fs/promises'; +import * as path from 'path'; +import * as os from 'os'; +import { startWarmupCommand } from './warmup'; + +describe('startWarmupCommand', () => { + test('waits for the command before reporting warmup complete', async () => { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'warmup-')); + const marker = path.join(tmp, 'warmed'); + try { + expect(await startWarmupCommand(`touch '${marker}'`, 5_000)).toBe('completed'); + expect(await fsp.readFile(marker, 'utf8')).toBe(''); + } finally { + await fsp.rm(tmp, { recursive: true, force: true }); + } + }); + + test('is a no-op when the command is unset or blank', async () => { + expect(await startWarmupCommand(undefined)).toBe('skipped'); + expect(await startWarmupCommand('')).toBe('skipped'); + expect(await startWarmupCommand(' ')).toBe('skipped'); + }); + + test('a failing command is surfaced without throwing', async () => { + expect(await startWarmupCommand('/nonexistent-warmup-binary --flag')).toBe('failed'); + expect(await startWarmupCommand('exit 7')).toBe('failed'); + }); + + test('terminates a warmup that exceeds its startup budget', async () => { + const started = Date.now(); + expect(await startWarmupCommand('sleep 10', 30)).toBe('timed_out'); + expect(Date.now() - started).toBeLessThan(2_000); + }); +}); diff --git a/api/src/warmup.ts b/api/src/warmup.ts new file mode 100644 index 0000000..f66d3da --- /dev/null +++ b/api/src/warmup.ts @@ -0,0 +1,71 @@ +import { spawn } from 'child_process'; +import { logger } from './logger'; + +export type WarmupOutcome = 'skipped' | 'completed' | 'failed' | 'timed_out'; + +/** + * Bounded warmup command run once at sandbox-API startup, OUTSIDE any job + * sandbox and before the HTTP listener becomes ready. + * + * Purpose: pre-fault large read-mostly assets into the VM's page cache before + * the first real job needs them. MicroVM rootfs reads are lazy and slow on + * first touch, so a heavyweight import (e.g. chdb's ~400MB shared object) can + * take 30-120s cold but milliseconds warm — a boot-time + * `SANDBOX_WARMUP_COMMAND="/pkgs/python/3.14.4/bin/python3 -c 'import chdb'"` + * moves that cost into the boot window instead of racing the user's first job. + * + * Failures stay best-effort (logged, then startup continues), but completion is + * awaited and bounded. This is important for Lambda MicroVM images: snapshotting + * a still-running detached warmup would clone that process into every VM and + * would not guarantee warm state before the endpoint became reachable. + */ +export async function startWarmupCommand( + command = process.env.SANDBOX_WARMUP_COMMAND, + timeoutMs = Number(process.env.SANDBOX_WARMUP_TIMEOUT_MS) || 180_000, +): Promise { + if (command == null || command.trim() === '') { + return 'skipped'; + } + const boundedTimeoutMs = Number.isFinite(timeoutMs) && timeoutMs > 0 + ? Math.floor(timeoutMs) + : 180_000; + const startedAt = Date.now(); + return new Promise((resolve) => { + let settled = false; + let timer: ReturnType | undefined; + const finish = (outcome: WarmupOutcome, code?: number | null, err?: unknown): void => { + if (settled) return; + settled = true; + if (timer) clearTimeout(timer); + const details = { command, code, elapsedMs: Date.now() - startedAt, err }; + if (outcome === 'completed') logger.info(details, 'Sandbox warmup command finished'); + else logger.warn(details, `Sandbox warmup command ${outcome.replace('_', ' ')}`); + resolve(outcome); + }; + let child: ReturnType; + try { + child = spawn('bash', ['-c', command], { + detached: true, + stdio: 'ignore', + }); + } catch (err) { + finish('failed', undefined, err); + return; + } + child.once('exit', code => finish(code === 0 ? 'completed' : 'failed', code)); + child.once('error', err => finish('failed', undefined, err)); + timer = setTimeout(() => { + try { + if (child.pid != null && process.platform !== 'win32') { + process.kill(-child.pid, 'SIGKILL'); + } else { + child.kill('SIGKILL'); + } + } catch { + child.kill('SIGKILL'); + } + finish('timed_out'); + }, boundedTimeoutMs); + logger.info({ command, timeoutMs: boundedTimeoutMs }, 'Sandbox warmup command started'); + }); +} diff --git a/api/src/workspace-isolation.test.ts b/api/src/workspace-isolation.test.ts index a2167bd..0888fbd 100644 --- a/api/src/workspace-isolation.test.ts +++ b/api/src/workspace-isolation.test.ts @@ -5,6 +5,7 @@ import * as path from 'path'; import { config } from './config'; import { SANDBOX_WORKSPACE_MODE, + SESSION_WORKSPACE_ID, SandboxJobUidPool, applySandboxPathPermissionsNoFollow, assertWorkspaceOwnershipCapability, @@ -14,9 +15,11 @@ import { compatibilityModeForSkippedChown, createSandboxWorkspace, createWorkspaceId, + ensureSessionWorkspace, prepareWorkspaceRoot, quarantineModeForUid, reapStaleWorkspaces, + resetSessionWorkspace, retainWorkspaceCleanupUntilRemoved, retainedWorkspaceCleanupCount, retryRetainedWorkspaceCleanups, @@ -24,6 +27,8 @@ import { workspaceOwnershipCapabilityErrors, } from './workspace-isolation'; +const nonRootIdentity = { slot: 0, uid: 65534, gid: 65534, perJobUid: false } as const; + const tmpRoots: string[] = []; const savedPerJobUids = config.per_job_uids; @@ -101,6 +106,46 @@ describe('sandbox UID slot pool', () => { }); }); +describe('persistent session workspace', () => { + test('ensureSessionWorkspace uses a stable id and preserves contents across calls', async () => { + config.per_job_uids = false; + const root = await mkroot('session-ws-'); + const first = await ensureSessionWorkspace(nonRootIdentity, root); + expect(first.workspaceId).toBe(SESSION_WORKSPACE_ID); + expect(first.dir).toBe(path.join(root, SESSION_WORKSPACE_ID)); + + await fsp.writeFile(path.join(first.dir, 'state.txt'), 'kept'); + const second = await ensureSessionWorkspace(nonRootIdentity, root); + expect(second.dir).toBe(first.dir); + expect(await fsp.readFile(path.join(second.dir, 'state.txt'), 'utf8')).toBe('kept'); + }); + + test('the reaper never removes the active session workspace', async () => { + config.per_job_uids = false; + const root = await mkroot('session-ws-reaper-'); + const lease = await ensureSessionWorkspace(nonRootIdentity, root); + await fsp.writeFile(path.join(lease.dir, 'keep.txt'), 'x'); + const removed = await reapStaleWorkspaces({ root, removeAll: true }); + expect(removed).toBe(0); + expect(await fsp.readFile(path.join(lease.dir, 'keep.txt'), 'utf8')).toBe('x'); + }); + + test('resetSessionWorkspace wipes the directory and unprotects it from the reaper', async () => { + config.per_job_uids = false; + const root = await mkroot('session-ws-reset-'); + const lease = await ensureSessionWorkspace(nonRootIdentity, root); + await fsp.writeFile(path.join(lease.dir, 'gone.txt'), 'x'); + + expect(await resetSessionWorkspace(root)).toBe(true); + await expect(fsp.access(lease.dir)).rejects.toBeDefined(); + + /* After reset a leftover dir with the same name is reapable again. */ + await fsp.mkdir(lease.dir, { recursive: true }); + const removed = await reapStaleWorkspaces({ root, removeAll: true }); + expect(removed).toBe(1); + }); +}); + describe('sandbox workspace root and reaper', () => { test('prepares a non-symlink workspace root with non-listable traversal mode', async () => { config.per_job_uids = false; diff --git a/api/src/workspace-isolation.ts b/api/src/workspace-isolation.ts index 1c0faa8..d5a0aa2 100644 --- a/api/src/workspace-isolation.ts +++ b/api/src/workspace-isolation.ts @@ -5,6 +5,7 @@ import * as path from 'path'; import type { Dirent } from 'fs'; import { config } from './config'; import { logger } from './logger'; +import { SANDBOX_DIR_MODE } from './validation'; export const SANDBOX_WORKSPACE_ROOT = '/tmp/sandbox'; /* NsJail's bind-mount setup needs execute permission on the source path and @@ -16,6 +17,12 @@ export const SANDBOX_READONLY_FILE_MODE = 0o444; export const SANDBOX_INSIDE_UID = 65534; export const SANDBOX_INSIDE_GID = 65534; export const WORKSPACE_ID_PREFIX = 'ws_'; +/* Fixed id for the single persistent session workspace (stateful mode). One + * per runner process — a session VM runs exactly one runtime session, whose + * executions serialize on the control-plane lock. Distinct from the random + * `ws_` per-job ids and kept in `activeWorkspaceIds` so the reaper never + * removes it mid-session. */ +export const SESSION_WORKSPACE_ID = 'session'; const RETAINED_WORKSPACE_RETRY_MS = 5_000; export class SandboxWorkspaceIsolationError extends Error { @@ -360,6 +367,51 @@ export async function createSandboxWorkspace( throw new SandboxWorkspaceIsolationError('Unable to allocate a unique sandbox workspace ID'); } +/** + * Idempotently ensures the single persistent session workspace exists and is + * owned by `identity`. Unlike `createSandboxWorkspace`, the directory is + * stable (`SESSION_WORKSPACE_ID`) and its contents are preserved across calls + * — reused by every execution in a stateful session. Registered in + * `activeWorkspaceIds` so the stale-workspace reaper skips it. + */ +export async function ensureSessionWorkspace( + identity: SandboxJobIdentity, + root = SANDBOX_WORKSPACE_ROOT, +): Promise { + if (identity.perJobUid) assertWorkspaceOwnershipCapability(); + await prepareWorkspaceRoot(root); + const dir = path.join(root, SESSION_WORKSPACE_ID); + try { + await fsp.mkdir(dir, { mode: SANDBOX_WORKSPACE_MODE }); + await applySandboxPathPermissions(dir, identity, SANDBOX_WORKSPACE_MODE); + } catch (error) { + if (errorCode(error) !== 'EEXIST') { + try { await fsp.rm(dir, { recursive: true, force: true }); } catch { /* best effort */ } + throw error; + } + /* Already present from a prior job in this session — keep contents, + * only (re)assert ownership so the pinned identity can read/write. */ + await applySandboxPathPermissions(dir, identity, SANDBOX_WORKSPACE_MODE); + } + activeWorkspaceIds.add(SESSION_WORKSPACE_ID); + return { workspaceId: SESSION_WORKSPACE_ID, dir, identity }; +} + +/** Tears down the persistent session workspace (session reset / VM terminate). */ +export async function resetSessionWorkspace(root = SANDBOX_WORKSPACE_ROOT): Promise { + const dir = path.join(root, SESSION_WORKSPACE_ID); + try { + await fsp.rm(dir, { recursive: true, force: true }); + return true; + } catch (error) { + logger.error({ dir, err: error }, 'Failed to reset session workspace'); + await quarantineWorkspace(dir); + return false; + } finally { + activeWorkspaceIds.delete(SESSION_WORKSPACE_ID); + } +} + async function quarantineWorkspace(dir: string): Promise { try { const st = await fsp.lstat(dir); @@ -530,3 +582,58 @@ export function fallbackSandboxIdentity(): SandboxJobIdentity { perJobUid: false, }; } + +/** + * Creates every directory from `root` (exclusive) to `target` (inclusive) one + * component at a time, refusing to descend through a symlink. + * `fsp.mkdir(dir, { recursive: true })` would instead follow a symlinked + * ancestor a prior session turn planted and create directories OUTSIDE the + * workspace as root — so any privileged write into a PERSISTENT session + * workspace (input priming, pushed file delivery) must build the path with + * no-follow semantics first. Fresh per-job workspaces can never contain such a + * link, so callers may skip it there. + * + * Pass `identity` to give directories this call creates the sandbox dir mode + * and ownership; without it they inherit the runner's (root) ownership, which + * would leave the sandbox UID unable to write inside them. + */ +export async function ensureDirNoFollow( + root: string, + target: string, + identity?: SandboxJobIdentity, +): Promise { + const rel = path.relative(root, target); + if (!rel || rel === '.') return; + if (rel === '..' || rel.startsWith('..' + path.sep)) { + throw new SandboxWorkspaceIsolationError(`Workspace path escapes the sandbox: ${target}`); + } + let cursor = root; + for (const part of rel.split(path.sep).filter(Boolean)) { + cursor = path.join(cursor, part); + const st = await fsp.lstat(cursor).catch(() => null); + if (st == null) { + /* Parent is already validated as a real dir; create just this component + * (non-recursive, so it can't follow a link). Tolerate a concurrent + * sibling prime having just created it. */ + let created = true; + await fsp.mkdir(cursor).catch((err: NodeJS.ErrnoException) => { + if (err.code !== 'EEXIST') throw err; + created = false; + }); + if (created && identity) { + await applySandboxPathPermissions(cursor, identity, SANDBOX_DIR_MODE); + } + continue; + } + if (st.isSymbolicLink()) { + throw new SandboxWorkspaceIsolationError( + `Refusing to write through symlinked workspace path: ${path.relative(root, cursor)}`, + ); + } + if (!st.isDirectory()) { + throw new SandboxWorkspaceIsolationError( + `Workspace ancestor is not a directory: ${path.relative(root, cursor)}`, + ); + } + } +} diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md new file mode 100644 index 0000000..6b78315 --- /dev/null +++ b/docs/lambda-microvm/README.md @@ -0,0 +1,586 @@ +# Stateful Code Sessions on AWS Lambda MicroVMs + +This directory documents and provisions the **optional** AWS Lambda MicroVM +execution backend for the CodeAPI sandbox. It turns the semi-stateless Code +Interpreter into one that offers **perceived-indefinite stateful sessions**: a +warm per-session workspace plus checkpoint/restore across the VM's 8-hour +lifetime, without changing the default HTTP behavior. + +- Config reference and knobs → below. +- One-command AWS prerequisites → [`terraform/`](./terraform). +- Image build helper → [`../../service/scripts/create-microvm-image.ts`](../../service/scripts/create-microvm-image.ts). + +> This is a config-gated feature. With `CODEAPI_SANDBOX_BACKEND` unset, nothing +> here is active and the sandbox behaves exactly as before. + +--- + +## The cross-repo picture + +Stateful sessions span three repos. Each owns one layer, and they degrade +gracefully out of order (the wire fields are additive and ignored when absent). + +| Repo | Provides | Key artifact | +|---|---|---| +| **code-interpreter** (this repo) | The Code API service + the Lambda MicroVM backend, the runner's persistent workspace, checkpoint/restore, and the session registry. Owns **all** AWS config. | `CODEAPI_SANDBOX_BACKEND=lambda-microvm` | +| **@librechat/agents** | The SDK surface: `toolExecution.sandbox.statefulSessions` and the `statefulSessions` tool-factory param. Stamps a per-conversation session hint on `/exec`. | `runtime_session_hint` on the wire | +| **LibreChat** | The `stateful_code_sessions` app capability + a per-agent Agent Builder toggle, wired into `createRun`. | endpoints.agents capability + agent toggle | + +**Trust boundary:** LibreChat and the agents SDK never learn any AWS +configuration. They speak the same `/exec` HTTP protocol as always, plus one +optional `runtime_session_hint` field. Everything AWS — backend selection, the +image ARN, roles, connectors, the checkpoint bucket, credentials — lives only in +this service's environment. An operator can switch this service between `http` +and `lambda-microvm` (or run with no AWS at all) with zero changes upstream. + +--- + +## How a request flows + +``` +agent tool call + │ POST /exec (+ optional runtime_session_hint) + ▼ +CodeAPI service ── derive runtime_session_id = hash(tenant, user, hint) + │ │ + │ Redis session registry (SET NX lock, generation fence) + ▼ │ +SandboxBackend (lambda-microvm) + │ find-or-launch ONE MicroVM per runtime_session_id + ▼ +RunMicrovm ─► warm VM ─► CreateMicrovmAuthToken ─► POST /api/v2/execute + │ (X-Runtime-Session-Id header = session mode on) + ▼ +runner reuses ONE /mnt/data workspace across calls + │ post-exec, lock held: checkpoint /mnt/data ──► S3 + ▼ +restore-on-relaunch: a replacement VM pulls the S3 checkpoint before first exec +``` + +Two independent planes: **presentation/orchestration** (the registry + backend, +which own identity and durability) and **compute** (the MicroVM, which is a cheap, +disposable, resumable cache). If the VM dies, its replacement resumes from the +last successfully committed checkpoint; work after a skipped or failed +checkpoint may need to be run again. + +--- + +## Prerequisites + +- An AWS account with **Lambda MicroVMs available** in your region (a new + service; confirm regional availability first). No default region is assumed — + every call passes one explicitly. +- **AWS CLI ≥ 2.35** if you want to poke the `aws lambda-microvms` CLI directly + (older CLIs lack the commands). The scripted image build uses the JS SDK and + doesn't need a recent CLI. +- Terraform ≥ 1.9 for the prerequisites module. +- Docker with `buildx` (arm64) to build the runner image. +- Redis (the CodeAPI service already depends on it) for the session registry. +- An S3-compatible store for checkpoints (real S3 in prod; MinIO for local dev). +- For the hardened path below, an internal egress-gateway deployment plus a VPC + network connector/security group that permits the MicroVM to reach only that + gateway. + +--- + +## Hardened deployment runbook + +### 1. Provision AWS prerequisites (Terraform) + +`terraform.tfvars.example` targets the disposable AIML-dev stack and explicitly +enables destructive teardown. Remove its three `*_force_*` overrides before +applying it to a retained environment; the module defaults them to `false`. + +```bash +cd docs/lambda-microvm/terraform +cp terraform.tfvars.example terraform.tfvars # edit region, name_prefix, image_name +terraform init -lockfile=readonly && terraform apply +``` + +This creates the private runner ECR repository, checkpoint bucket (encrypted, +versioned, lifecycle-expired), artifact bucket, the **build role** (trust +includes `sts:TagSession`; permissions include scoped logs, artifact read, and +ECR pull), a **logging-only execution role**, the worker MicroVM/checkpoint +policies, and the build + runtime CloudWatch log groups. Set +`codeapi_worker_role_name` to attach both worker policies directly, or attach the +two output policy ARNs in the stack that owns the worker role. Capture the +outputs: + +```bash +terraform output +# runner_ecr_repository_url, build_role_arn, execution_role_arn, +# worker_microvm_control_policy_arn, checkpoint_access_policy_arn, ... +``` + +The checked-in provider lock file makes `terraform init` reproducible. + +### 2. Build + push the runner image, upload the code-artifact + +```bash +cd ../../.. # repo root +export AWS_PROFILE=... AWS_REGION=us-east-1 +export ECR_URI=$(terraform -chdir=docs/lambda-microvm/terraform output -raw runner_ecr_repository_url) +export S3_URI=s3://$(terraform -chdir=docs/lambda-microvm/terraform output -raw artifact_bucket)/runner +export IMAGE_TAG=$(git rev-parse --short HEAD) +scripts/build-lambda-microvm-artifact.sh build push zip upload +# → uploads s3:///runner/runner-.zip +``` + +The runner image target is `lambda-microvm-runner` in `api/Dockerfile` +(`FROM sandbox-build` + `/pkgs`, `PORT=8080`, `SANDBOX_SESSION_WORKSPACE_ENABLED=true`). +The managed ECR repository uses immutable tags. If a push already succeeded, +choose a fresh `IMAGE_TAG` (for example, append the CI run ID) instead of +rerunning a push to the same commit tag. The `push` stage captures the resulting +ECR digest and the `zip` stage writes `FROM @sha256:...`; it refuses +to produce a mutable tag-based artifact. The `upload` stage additionally +verifies the artifact's repository, tag, digest, and SHA-256 before writing to +S3. When running `zip` and `upload` in a separate clean workspace, pass the +digest explicitly in the same invocation: + +```bash +IMAGE_DIGEST=sha256:<64-hex-digest> \ + scripts/build-lambda-microvm-artifact.sh zip upload +``` + +### 3. Generate the split execution-manifest keys + +The worker signs each execution manifest; the runner only receives the public +verifier. Never bake the private key into the image. + +```bash +MANIFEST_KEY_DIR="$(mktemp -d)" +trap 'rm -rf "$MANIFEST_KEY_DIR"' EXIT +ORIGINAL_UMASK="$(umask)" +umask 077 +openssl genpkey -algorithm ED25519 -out "$MANIFEST_KEY_DIR/private.pem" +export CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY=$( + openssl pkey -in "$MANIFEST_KEY_DIR/private.pem" -outform DER | base64 | tr -d '\n' +) +export SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY=$( + openssl pkey -in "$MANIFEST_KEY_DIR/private.pem" -pubout -outform DER | base64 | tr -d '\n' +) +rm -rf "$MANIFEST_KEY_DIR" +trap - EXIT +umask "$ORIGINAL_UMASK" +unset MANIFEST_KEY_DIR ORIGINAL_UMASK +``` + +Store the private value in the worker's secret manager. Put only the public +value into the runner image environment below, then unset the private shell +variable after the secret is stored. No private-key file is left in the checkout. + +### 4. Create the MicroVM image (hookless) + +This walkthrough uses the hardened egress path so output uploads and tool calls +cannot bypass policy. Deploy the `egress-gateway` target from +`service/Dockerfile` behind an internal URL before accepting traffic. Configure +that gateway workload (and no API/worker workload) with: + +```bash +CODEAPI_HARDENED_SANDBOX_MODE=true +CODEAPI_EGRESS_GRANT_SECRET= +CODEAPI_INTERNAL_SERVICE_TOKEN= +EGRESS_GATEWAY_FILE_SERVER_URL=http://file-server.internal:3000 +EGRESS_GATEWAY_TOOL_CALL_SERVER_URL=http://tool-call-server.internal:3033 +EGRESS_GATEWAY_PORT=3190 +CODEAPI_EGRESS_LEDGER_REQUIRED=true +REDIS_HOST= +# REDIS_PORT / REDIS_USERNAME / REDIS_PASSWORD / TLS settings as required +``` + +The gateway URL must be reachable by both the worker and the MicroVM VPC +connector. Its security group should accept the MicroVM connector on only the +gateway port. The gateway's file/tool upstreams stay private. + +```bash +cd service +export EGRESS_GATEWAY_URL=https://egress.internal +export MICROVM_IMAGE_ENV_JSON="$( + bun -e 'process.stdout.write(JSON.stringify({ + CODEAPI_HARDENED_SANDBOX_MODE: "true", + EGRESS_GATEWAY_URL: process.env.EGRESS_GATEWAY_URL, + SANDBOX_CHECKPOINT_MAX_BYTES: "536870912", + SANDBOX_ALLOWED_LOCAL_NETWORK_PORT: "3033", + SANDBOX_FORWARD_TARGET: process.env.EGRESS_GATEWAY_URL, + SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY: process.env.SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY, + SANDBOX_REQUIRE_EGRESS_MANIFEST: "true" + }))' +)" +AWS_PROFILE=... bun scripts/create-microvm-image.ts \ + --name codeapi-session \ + --artifact s3:///runner/runner-.zip \ + --build-role $(terraform -chdir=../docs/lambda-microvm/terraform output -raw build_role_arn) \ + --region us-east-1 \ + --env-json "$MICROVM_IMAGE_ENV_JSON" +# Optional for a fully repeatable build: +# --base-version +# → prints the exact LAMBDA_MICROVM_IMAGE_ARN and LAMBDA_MICROVM_IMAGE_VERSION +``` + +The helper builds hookless with `additionalOsCapabilities:["ALL"]` and +`SANDBOX_USE_CGROUPV2=false` baked in — the working config (see +[Runbook gotchas](#runbook-gotchas)). **Runner env is baked at image-build +time** (RunMicrovm does not inject it later), so pass your deployment's +egress / manifest config via `--env-json` (or `MICROVM_IMAGE_ENV_JSON`) — +for hardened mode that means `CODEAPI_HARDENED_SANDBOX_MODE`, +`EGRESS_GATEWAY_URL`, +`SANDBOX_CHECKPOINT_MAX_BYTES`, `SANDBOX_ALLOWED_LOCAL_NETWORK_PORT`, +`SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY`, `SANDBOX_REQUIRE_EGRESS_MANIFEST`, and +a `SANDBOX_FORWARD_TARGET` whose host and port match the gateway URL. Never bake +`CODEAPI_INTERNAL_SERVICE_TOKEN`, +`CODEAPI_EGRESS_GRANT_SECRET`, a private manifest key, or direct file/tool +service URLs into the runner. The worker fetches by-reference input objects and +pushes an authenticated cache batch into the runner; the runner does not receive +raw file-server object IDs or need `FILE_SERVER_URL` for inputs. To ship new +runner code later, re-run +`build … push zip upload`, call the helper with `--update`, and deploy the newly +printed pinned version before replacing workers. + +`--base-version` (or `MICROVM_BASE_IMAGE_VERSION`) pins the AWS-managed +`--base-image` version as well as the CodeAPI artifact. If omitted, the helper +prints that it used the managed base's latest version; record that fact for a +dev experiment, but pin an immutable base version for reproducible release +builds. + +### 5. Configure the CodeAPI service + +```bash +CODEAPI_SANDBOX_BACKEND=lambda-microvm +CODEAPI_RUNTIME_SESSION_MODE=affinity # warm sessions + checkpoints +LAMBDA_MICROVM_IMAGE_ARN= +LAMBDA_MICROVM_IMAGE_VERSION= # required for affinity/strict +LAMBDA_MICROVM_EXECUTION_ROLE_ARN= +LAMBDA_MICROVM_LOG_GROUP= # both this AND the role are needed for VM stdout +LAMBDA_MICROVM_REGION=us-east-1 +LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS=arn:aws:lambda:us-east-1:aws:network-connector:aws-network-connector:ALL_INGRESS +LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS= + +# set on BOTH API and worker; use the same URL/token as the gateway deployment +CODEAPI_HARDENED_SANDBOX_MODE=true +EGRESS_GATEWAY_URL=https://egress.internal +CODEAPI_INTERNAL_SERVICE_TOKEN= + +# worker only: source for authorized by-reference input objects. The worker +# pushes bytes into the runner cache; never bake this URL into the runner image. +FILE_SERVER_URL=http://file-server.internal:3000 + +# checkpoints (S3-compatible, same client as file-server) +CODEAPI_CHECKPOINT_BUCKET= +CODEAPI_CHECKPOINT_MAX_BYTES=536870912 +MINIO_ENDPOINT=https://s3.us-east-1.amazonaws.com +# MINIO_PORT=443 # optional explicit override; otherwise the URL scheme supplies 443 +MINIO_REGION=us-east-1 +# ECS task role, EC2 instance profile, and IRSA/web identity are loaded +# automatically. Only non-role/local deployments need a complete static pair: +# MINIO_ACCESS_KEY=... +# MINIO_SECRET_KEY=... + +# worker-only signer; do not bake this into the runner image +CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY= +``` + +Do not set `CODEAPI_EGRESS_GRANT_SECRET` on the API or worker; the hardened +startup checks reject it there. The worker authenticates to the gateway with +`CODEAPI_INTERNAL_SERVICE_TOKEN`, the gateway mints the per-execution grant, and +the runner receives only that short-lived grant plus the signed manifest. +Terraform emits the worker IAM policies but intentionally does not create your +VPC connector, gateway load balancer, or security groups; those must be owned by +the deployment's network stack. + +`PTC_MODE` must be `replay` (the default) or unset — see +[Programmatic Tool Calling](#programmatic-tool-calling-ptc). + +### 6. Verify + +Enable the capability + per-agent toggle in LibreChat and run a two-message +conversation: write `42` to `/mnt/data/answer.txt`, then read it back in a +follow-up message. With the session backend it reads `42`; with the toggle off, +the follow-up sees no file. (The LibreChat PR documents the full acceptance test +and the no-infra wiring smoke.) + +--- + +## Configuration reference + +Worker/service names appear in `service/src/config.ts`; runner-image names +appear in `api/src/config.ts`. + +### Backend selection + +| Env | Default | Meaning | +|---|---|---| +| `CODEAPI_SANDBOX_BACKEND` | `http` | `http` (byte-identical to today) or `lambda-microvm`. | +| `CODEAPI_RUNTIME_SESSION_MODE` | `stateless` | `stateless` \| `affinity` \| `strict`. See [Operating modes](#operating-modes). | +| `CODEAPI_RUNTIME_SESSION_LOCK_WAIT_MS` | `15000` | How long a stateful execution waits for the session lock before returning `RUNTIME_SESSION_BUSY` (HTTP 409). | + +### MicroVM launch + +| Env | Default | Meaning | +|---|---|---| +| `LAMBDA_MICROVM_IMAGE_ARN` | — (required) | The image created in step 4. | +| `LAMBDA_MICROVM_IMAGE_VERSION` | latest in stateless mode; required in affinity/strict | Exact image version. Stateful startup rejects an unpinned version so a rolling image update cannot silently mix workspace sessions across revisions. | +| `LAMBDA_MICROVM_EXECUTION_ROLE_ARN` | — | Logging-only role. Required (with the log group below) for runtime VM stdout to reach CloudWatch. | +| `LAMBDA_MICROVM_LOG_GROUP` | — | CloudWatch log group sent on `RunMicrovm`. Needed alongside the execution role or stdout goes nowhere. | +| `LAMBDA_MICROVM_REGION` | SDK default | Region for the lambda-microvms client. | +| `LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS` | — | Comma-separated. Inbound HTTPS to the VM. | +| `LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS` | — | Comma-separated. Outbound from the VM. Required in hardened mode. | +| `LAMBDA_MICROVM_PORT` | `8080` | Runner port. | +| `LAMBDA_MICROVM_MAX_DURATION_SECONDS` | `28800` | Hard lifetime ceiling (≤ 8h). | +| `LAMBDA_MICROVM_IDLE_SECONDS` | `1800` | idlePolicy: auto-suspend after this idle. | +| `LAMBDA_MICROVM_SUSPEND_SECONDS` | `1800` | idlePolicy: auto-terminate after this suspended. | +| `LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS` | `300` | Requested AWS proxy-token lifetime. Tokens are minted afresh for each probe, push, checkpoint, restore, or execute request; the multi-probe launch-readiness loop refreshes its token near expiry. | +| `LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS` | `60000` | Budget for RunMicrovm → RUNNING. | +| `LAMBDA_MICROVM_HEALTH_TIMEOUT_MS` | `5000` | Health check budget. | +| `LAMBDA_MICROVM_LAUNCH_TPS` | `4` | Client-side launch throttle (headroom under AWS's 5 TPS cap). Suspend/resume are driven by the configured AWS idle policy, not worker control-plane calls. | +| `LAMBDA_MICROVM_TOKEN_TPS` | `8` | Shared worker throttle for fresh `CreateMicrovmAuthToken` calls. | +| `LAMBDA_MICROVM_ALLOW_SHELL` | `false` | Must stay false in prod (shell auth token → IAM-deny). | + +### Hardened egress split + +| Env | Placement | Meaning | +|---|---|---| +| `CODEAPI_HARDENED_SANDBOX_MODE=true` | API, worker, gateway, runner image | Enables each process's fail-closed startup checks. Because runner env is immutable, it must be present during image creation. | +| `EGRESS_GATEWAY_URL` | API, worker, runner image | Bare gateway origin (no path/query). Worker creates and restores grants; runner routes file/tool egress through it. | +| `CODEAPI_INTERNAL_SERVICE_TOKEN` | API, worker, gateway, file server, tool-call server; **never runner** | Authenticates internal control/upstream calls. Use one strong secret value across these workloads. | +| `CODEAPI_EGRESS_GRANT_SECRET` | Gateway only | Seals per-execution grants. Hardened API/worker startup rejects this secret. | +| `EGRESS_GATEWAY_FILE_SERVER_URL` / `_TOOL_CALL_SERVER_URL` | Gateway only | Private upstream origins. Do not expose them to the runner. | +| `FILE_SERVER_URL` | Worker only | Direct source used to fetch authorized by-reference inputs before pushing them into the runner cache. Do not bake it into the runner. | +| `CODEAPI_EGRESS_LEDGER_REQUIRED=true` / `REDIS_*` | Gateway | Makes grant replay/revocation state fail closed in Redis. | +| `SANDBOX_ALLOWED_LOCAL_NETWORK_PORT` / `SANDBOX_FORWARD_TARGET` | Runner image | Configures the narrow tool-call socket proxy; it does not start or expose the socket by itself. The target host/port must match `EGRESS_GATEWAY_URL`. | +| `CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY` | Worker only | Signs the manifest after the gateway returns the scoped grant. | +| `SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY` / `SANDBOX_REQUIRE_EGRESS_MANIFEST=true` | Runner image | Verifies and requires the signed per-execution manifest. | + +The Unix socket is mounted only for an execution whose body-bound, signed +manifest carries `tool_call_socket=true`; the service's matching request flag is +validated against that scope. Ordinary execution and replay PTC leave the claim +false. This per-execution capability is internal control-plane data, not an +operator environment switch. + +### Checkpoints (affinity/strict only) + +| Env | Default | Meaning | +|---|---|---| +| `CODEAPI_SESSION_CHECKPOINTS` | `true` | `false` disables checkpoint/restore. Warm reuse still works, but a VM replacement starts clean and only receives the inputs declared by that request. | +| `CODEAPI_CHECKPOINT_BUCKET` | `MINIO_BUCKET` | Checkpoint bucket. | +| `CODEAPI_CHECKPOINT_PREFIX` | `rtsx-checkpoints/` | Key prefix. Immutable objects are `/<20-digit-sequence>.tar.gz`; a sibling `.committed` marker makes a fenced upload eligible for Redis-loss recovery. | +| `CODEAPI_CHECKPOINT_MAX_BYTES` | `536870912` | Max checkpoint size (512 MiB). | +| `SANDBOX_CHECKPOINT_MAX_BYTES` | `536870912` | Runner-side ceiling for streamed checkpoint restores. Bake it into the MicroVM image and keep it equal to worker `CODEAPI_CHECKPOINT_MAX_BYTES`. | +| `CODEAPI_CHECKPOINT_TIMEOUT_MS` | `60000` | Checkpoint transfer budget. | +| `MINIO_ENDPOINT` / `_PORT` / `_USE_SSL` / `_REGION` | — | S3-compatible endpoint configuration. Absolute URLs use their scheme-default port; `MINIO_PORT` explicitly overrides it. Bare MinIO-style endpoints default to port 9000 and use `_USE_SSL` to select the scheme. Point at real S3 in prod. | +| `MINIO_ACCESS_KEY` / `_SECRET_KEY` / `_SESSION_TOKEN` | workload IAM provider | Optional complete static credential set for local/non-role deployments. ECS, EC2, and web-identity/IRSA credentials are loaded automatically when the pair is absent. | + +The worker rejects oversized checkpoint objects before transfer and the runner +independently bounds the streamed restore body. Keep +`CODEAPI_CHECKPOINT_MAX_BYTES` and `SANDBOX_CHECKPOINT_MAX_BYTES` aligned; a +smaller runner limit makes otherwise valid worker checkpoints unrestorable, +while a larger runner limit weakens the receiver-side fail-closed ceiling. + +--- + +## Operating modes + +`CODEAPI_RUNTIME_SESSION_MODE` picks the tradeoff: + +- **`stateless`** — no registry. One VM per execution: run → execute → + terminate. MicroVM isolation per call, but no warm sessions and no + checkpoints. Correct and simple; the safest first AWS step. +- **`affinity`** — find-or-launch one warm VM per `runtime_session_id`. Lock + contention past `LOCK_WAIT_MS` returns HTTP 409; it never executes cold and + silently loses session workspace state. Requests without a session hint still + use the stateless path. +- **`strict`** — same serialized stateful behavior, but requests without a + session hint are rejected instead of degrading to stateless. + +--- + +## Alternative AWS methods + +You do not have to adopt the whole stack at once. The knobs compose: + +**No AWS at all.** Leave `CODEAPI_SANDBOX_BACKEND` unset (`http`). Today's +behavior, no MicroVMs, no changes needed anywhere. + +**MicroVM isolation without sessions.** `lambda-microvm` + `stateless`. Every +execution gets a fresh, strongly-isolated Firecracker VM. No registry, no +checkpoints, no session workspace. Simplest way to get the isolation boundary. + +**Base container image and snapshot boundary.** The default runner uses a stock +`oven/bun` base and is **hookless** — session mode arrives per request via the +`X-Runtime-Session-Id` header, so no lifecycle hooks are needed. Lambda image +creation snapshots running processes, memory, and file descriptors. The +Node-based tool-call proxy is therefore not started by the entrypoint; it starts +and becomes ready lazily only after a post-restore, authenticated execution is +explicitly authorized for the tool-call socket. This keeps Node's embedded +OpenSSL process state out of the image snapshot. + +If a future change must run an OpenSSL-using process before the snapshot, use +AWS's snapshot-safe OpenSSL libraries or rebase the +`lambda-microvm-runner` target on the snapshot-compatible Lambda base container +image (`public.ecr.aws/lambda/microvms:al2023-minimal`). That base also contains +the hook-routing components required for build/runtime lifecycle hooks. See +[AWS's snapshot guidance](https://docs.aws.amazon.com/lambda/latest/dg/microvms-images-snapshots.html) +and [Runbook gotchas](#runbook-gotchas). + +**Checkpoint store.** The checkpoint client is MinIO-compatible. For local dev, +point `MINIO_*` at a local MinIO. For prod, point it at real S3 (endpoint +`https://s3..amazonaws.com`; its scheme supplies port 443 unless +`MINIO_PORT` explicitly overrides it) and attach +`checkpoint_access_policy_arn` to the workload role. The client loads ECS task +role, EC2 instance-profile, and web-identity/IRSA credentials; a complete static +`MINIO_ACCESS_KEY`/`MINIO_SECRET_KEY` pair remains available for local or +non-role deployments. + +**Egress posture.** For dev, the Lambda-managed `INTERNET_EGRESS` connector gives +default public egress, but do not combine it with a runner that has +`CODEAPI_HARDENED_SANDBOX_MODE=true`. Hardened prod requires the complete split +configuration in steps 4–5: hardened runner image, manifest verifier, matching +local forward target, hardened API/worker, separately deployed gateway with its +grant secret and Redis ledger, shared internal-service token, and a VPC +connector/SG locked to only that gateway. Startup then requires +`LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS`. MicroVMs default to public egress, so +this gate is deliberate. + +**Throughput / quota.** `RunMicrovm` is capped account-wide (~5 TPS default), so +stateless cold throughput is ~4 exec/s fleet-wide until you request a quota +raise. Warm sessions (affinity) amortize this away for repeat calls in a +conversation. Treat a fresh account as canary-only. + +--- + +## Programmatic Tool Calling (PTC) + +- **Replay PTC works** and is the only supported PTC mode on this backend + (startup rejects `PTC_MODE=blocking`). Replay externalizes continuation state + in Redis, so each round is an independent `/exec` that can land on a fresh + one-shot VM and stay correct. +- **Blocking PTC is rejected** — it needs a live tool-call socket held open + through the auth proxy mid-execution, which fights the short VM lifecycle. +- **Replay PTC never requests the socket.** Its payload and signed manifest keep + `tool_call_socket=false`, so the post-restore Node proxy remains stopped. +- **PTC does not yet get warm sessions.** `/exec/programmatic` doesn't derive a + `runtime_session_id`, so PTC rounds always take the stateless path even when + the conversation's `execute_code` has a warm VM. Each replay round therefore + costs a VM launch. Binding PTC into the session VM is a planned follow-up (the + hint is already on the wire). + +--- + +## Runbook gotchas + +Each of these cost a silent or blind failure during bring-up: + +- **Build role trust** must include `sts:TagSession` alongside `sts:AssumeRole`, + and permissions must include log group/stream writes, `s3:GetObject`, and ECR + pull for a private base — missing any yields a `CREATE_FAILED` build with an + **empty** `stateReason`. (The Terraform module gets this right.) +- **Build logs** live at `/aws/lambda-microvms/` (hyphen), **not** the + AWS docs' `/aws/lambda/microvms/` (slash) — the docs are wrong; verified + empirically, the slash path does not exist in a live account. Do not "correct" + the Terraform log group or IAM policy to the docs' path or you lose the build + logs. +- **Runtime VM stdout** needs BOTH a `cloudWatch` logging config on RunMicrovm + AND an `executionRoleArn`, or it goes nowhere. Set + `LAMBDA_MICROVM_EXECUTION_ROLE_ARN`. +- **nsjail inside the guest** needs `additionalOsCapabilities:["ALL"]` (for the + `/proc` mount, else EPERM) and `SANDBOX_USE_CGROUPV2=false` (the app container + can't read the cgroup v2 subtree). Both are baked into the image helper. Under + ALL caps nsjail runs `no_pivotroot` — weaker in-guest isolation, acceptable + because the MicroVM is the real trust boundary (one VM per session). +- **NOFILE**: the AL2023 guest hard-caps `RLIMIT_NOFILE` at 1024, below the + runner's default; the entrypoint raises the hard limit to 65536. Docker masks + this locally. +- **Hooks never route on a stock container image.** Enabling any runtime hook + forces the `/ready` build hook, which never reaches a stock container's + listener, so the build fails at the ready timeout. Stay hookless (the default) + unless you rebase on the Lambda base container image. +- **Do not start the Node tool-call proxy in the entrypoint.** Lambda image + creation snapshots live process memory. The runner intentionally starts Node + only after restore and only for a body-bound, manifest-authorized + `tool_call_socket` execution, avoiding cloned OpenSSL process state. Audit any + new pre-snapshot OpenSSL consumer against AWS's snapshot guidance or use the + snapshot-safe base/libraries. + +--- + +## Rolling out a new runner image + +Image updates create a new immutable version. Build and upload a unique +`IMAGE_TAG`, call the helper with `--update`, then deploy both values it prints: + +```bash +cd service +AWS_PROFILE=... bun scripts/create-microvm-image.ts \ + --update \ + --name codeapi-session \ + --artifact s3:///runner/runner-.zip \ + --build-role \ + --region "$AWS_REGION" \ + --env-json '' +``` + +This release has a pre-release checkpoint compatibility boundary: the +runner-owned sidecar is now `codeapi.session-meta.v2`, which stores stable +input-cache digests instead of the v1 masked object handles. Before deploying +this branch to AIML-dev, drain and terminate every v1 warm session, then discard +or recycle its checkpoints and registry records. Do not run old and new runner +image versions concurrently; a rolling mixed-image deployment is unsupported +across this boundary. + +For later schema-compatible rollouts, update +`LAMBDA_MICROVM_IMAGE_VERSION` before replacing workers. The registry's launch +fingerprint includes the resolved image version and launch/security settings, +so a worker will not reuse a warm VM created under an older configuration. +Retain the prior image version until its MicroVMs have drained or been +terminated. + +--- + +## Teardown + +```bash +# Terminate only VMs launched from THIS image, then delete the image. +# The server-side image filter and AWS CLI auto-pagination keep this scoped in a +# shared account. +export AWS_PROFILE=... +export AWS_REGION=us-east-1 +export IMAGE_ARN="arn:aws:lambda:us-east-1::microvm-image:codeapi-session" + +for microvm_id in $(aws lambda-microvms list-microvms \ + --image-identifier "$IMAGE_ARN" \ + --region "$AWS_REGION" \ + --query "items[?state!='TERMINATED' && state!='TERMINATING'].microvmId" \ + --output text); do + aws lambda-microvms terminate-microvm \ + --microvm-identifier "$microvm_id" \ + --region "$AWS_REGION" +done + +# Wait at most 10 minutes. JSON output applies the JMESPath query once to the +# complete auto-paginated response; text output would return one count per page. +attempt=0 +while :; do + remaining=$(aws lambda-microvms list-microvms \ + --image-identifier "$IMAGE_ARN" \ + --region "$AWS_REGION" \ + --query "length(items[?state!='TERMINATED'])" \ + --output json) + [ "$remaining" = "0" ] && break + attempt=$((attempt + 1)) + if [ "$attempt" -ge 120 ]; then + echo "Timed out waiting for $remaining MicroVM(s) to terminate" >&2 + exit 1 + fi + sleep 5 +done + +aws lambda-microvms delete-microvm-image \ + --image-identifier "$IMAGE_ARN" \ + --region "$AWS_REGION" + +# then the prerequisites +terraform -chdir=docs/lambda-microvm/terraform destroy +``` + +MicroVM images are billed as stored snapshots; running VMs bill while RUNNING and +suspended VMs bill at a reduced rate, so terminate stray VMs before deleting the +image. diff --git a/docs/lambda-microvm/terraform/.gitignore b/docs/lambda-microvm/terraform/.gitignore new file mode 100644 index 0000000..0bec28e --- /dev/null +++ b/docs/lambda-microvm/terraform/.gitignore @@ -0,0 +1,5 @@ +.terraform/ +*.tfstate +*.tfstate.* +terraform.tfvars +crash.log diff --git a/docs/lambda-microvm/terraform/.terraform.lock.hcl b/docs/lambda-microvm/terraform/.terraform.lock.hcl new file mode 100644 index 0000000..7728a0a --- /dev/null +++ b/docs/lambda-microvm/terraform/.terraform.lock.hcl @@ -0,0 +1,28 @@ +# This file is maintained automatically by "terraform init". +# Manual edits may be lost in future updates. + +provider "registry.terraform.io/hashicorp/aws" { + version = "6.56.0" + constraints = ">= 5.40.0, < 7.0.0" + hashes = [ + "h1:MhideOrA+/8rovaA9BC23G4dDrFNZ0mfrRMNptZBIS0=", + "h1:gSTd4VOv0lEjCGP5deZ4hRMBZhIOUbtS4AlCML+3gIo=", + "h1:iWz9BgFQaDPA1ChVGYvgI3MlUnx3wSQCA2ggYqDPcz8=", + "zh:2b3fbb3bebcc663b85d5fd9bbc2d131ab89322d696ff5c6ac6b7ffb7b5fe92e7", + "zh:30e56ccc7f33a7778ab323a28fe893d8e9200dc5fb92ccb7023bee808db3c1b0", + "zh:67dca271bef16547ef8ab5a6349f9bce39d91d7c1ae3d8388ada687ca774ba44", + "zh:824c812695b14d2fddad5e22339d4520e16d9c875f5d5095f29003f49a6fd124", + "zh:8372b12e30078d1df8b52e1285fd0d9d35160a7d18c3b0211f55229e5a832fd3", + "zh:8922b45ab65e272e8951f70d890444ddb3441170d8fa6f51298f24f20d21943d", + "zh:8fc4fb94ac8547717128cbcf296ac0ba0918738c4882029cbba46bd4812eb5e4", + "zh:9b12af85486a96aedd8d7984b0ff811a4b42e3d88dad1a3fb4c0b580d04fa425", + "zh:a2efcd0174b79c1dffce48a5984f137ebc6f67d2c2ee966b47210e1faeb05cf2", + "zh:a4a50aa269a19c1d664b0dd59ee8047ed8a4a5b449f54733518309feccce1f43", + "zh:a830d80316b3c3127c9e0c77b9d4f50702c9d1a884c08973ac50b326d9ac734a", + "zh:d7e1b9bb2df3cdde8381ab101a807ae54e72c156d13a6b73dae2b922dca0c9f5", + "zh:d87c2a1cfc03b01cf13dff3700898ab990e1817326728824da8499dd0129da72", + "zh:fbb1848a597263c66dc6587ec8c3e0586e505e36f99d23752763534f62a3fef7", + "zh:fcb54d08eb3c763f01223e12b4eacbd018478e8e67c6b8611940564dfbea3d90", + "zh:ff6df70b2b2f75bd9000630f131c02e6bc2b9172cdb2547030f3fc019a3f3bc5", + ] +} diff --git a/docs/lambda-microvm/terraform/README.md b/docs/lambda-microvm/terraform/README.md new file mode 100644 index 0000000..f9863b1 --- /dev/null +++ b/docs/lambda-microvm/terraform/README.md @@ -0,0 +1,84 @@ +# Terraform: Lambda MicroVM prerequisites + +Provisions the static AWS resources the CodeAPI Lambda MicroVM backend needs. The +MicroVM image and running VMs themselves are **not** Terraform-managed — the +`lambda-microvms` service has no TF resource yet, so those are created with +[`service/scripts/create-microvm-image.ts`](../../../service/scripts/create-microvm-image.ts) +and by the backend at runtime. See [../README.md](../README.md) for the full +walkthrough. + +## What it creates + +- **Checkpoint S3 bucket** — encrypted (SSE-S3), versioned, public-access + blocked, with current and noncurrent versions lifecycle-expired separately + (`checkpoint_retention_days` and + `checkpoint_noncurrent_retention_days`). +- **Artifact S3 bucket** — for the code-artifact zip (optional; reuse an existing + one with `create_artifact_bucket = false` + `artifact_bucket_name`), with + lifecycle cleanup after `artifact_retention_days`. +- **Private ECR repository** — for the arm64 runner image (optional; reuse an + existing repository with `create_ecr_repository = false`), retaining the + newest `ecr_max_image_count` images. +- **Build role** — assumed by Lambda during `create/update-microvm-image`. Trust + includes `sts:TagSession`; permissions include writes to the exact build log + group, `s3:GetObject` on the artifact bucket, and (optional) scoped + private-ECR pull. Getting this wrong yields a build failure with an empty + `stateReason`. +- **Execution role** — logging-only least-privilege, for `RunMicrovm`. +- **Worker control policy** — `Run/Get/TerminateMicrovm`, + `CreateMicrovmAuthToken`, and the dependent `iam:PassRole` / + `lambda:PassNetworkConnector` permissions required by `RunMicrovm`. The + worker does not call or receive permission for `SuspendMicrovm` or + `ResumeMicrovm`; the configured AWS idle policy performs those transitions. +- **CloudWatch log groups** — build (`/aws/lambda-microvms/`) and + runtime. +- **Checkpoint access** — an IAM policy for task-role/instance-profile/IRSA + credentials (preferred), or an optional IAM user + access key + (`create_checkpoint_access_user = true`) for non-role deployments. + +## Usage + +The example targets disposable AIML-dev and explicitly enables destructive +teardown; change the three `*_force_*` values to `false` before applying it to a +retained environment. + +```bash +cp terraform.tfvars.example terraform.tfvars # edit +terraform init -lockfile=readonly +terraform apply +terraform output +``` + +## Notes + +- Commit `.terraform.lock.hcl` (this module does) and keep + `terraform init -lockfile=readonly` in automation so an AWS provider release + cannot silently change a previously reviewed plan. +- Set `image_name` to match the `--name` you pass to `create-microvm-image.ts`, + so the build log group is pre-created at the exact path Lambda writes to. +- Set `codeapi_worker_role_name` to attach the MicroVM and checkpoint policies + directly to an existing task role; otherwise attach both output policy ARNs + in the deployment that owns that role. +- `runner_ecr_repository_url` is the `ECR_URI` consumed by the build script. +- `private_ecr = false` is only for a code-artifact Dockerfile that already + references a public image. It makes `runner_ecr_repository_url` null; the + provided build/push helper and hardened runbook use private ECR. +- `create_checkpoint_access_user = true` exposes `checkpoint_access_key_id` and + the sensitive `checkpoint_secret_access_key` outputs — use as `MINIO_ACCESS_KEY` + / `MINIO_SECRET_KEY`. The checkpoint client now loads ECS/EC2/IRSA credentials, + so prefer a role. A `sensitive` output is still stored in plaintext Terraform + state; use an encrypted, access-controlled remote backend and never copy state + into logs or artifacts. +- Managed artifacts expire after `artifact_retention_days`; ECR keeps the newest + `ecr_max_image_count` images. Use a unique `IMAGE_TAG` for every push because + the ECR repository is immutable, and retain every version needed for rollback. +- `artifact_force_destroy`, `checkpoint_force_destroy`, and `ecr_force_delete` + default to `false`, so a destroy cannot silently empty buckets or the image + repository. The checked-in `terraform.tfvars.example` is specifically an + AIML-dev/disposable-stack example and explicitly opts all three into + destructive teardown; do not copy those overrides to retained environments. +- Current checkpoint versions expire after `checkpoint_retention_days`. + Noncurrent S3 versions expire independently after + `checkpoint_noncurrent_retention_days` (one day by default), so bucket + versioning does not unexpectedly retain replaced checkpoint data for another + full current-version window. diff --git a/docs/lambda-microvm/terraform/main.tf b/docs/lambda-microvm/terraform/main.tf new file mode 100644 index 0000000..4f89976 --- /dev/null +++ b/docs/lambda-microvm/terraform/main.tf @@ -0,0 +1,411 @@ +data "aws_caller_identity" "current" {} +data "aws_partition" "current" {} + +locals { + account_id = data.aws_caller_identity.current.account_id + artifact_bucket = var.create_artifact_bucket ? aws_s3_bucket.artifact[0].id : var.artifact_bucket_name + runner_repository_arn = ( + var.private_ecr && var.create_ecr_repository + ? aws_ecr_repository.runner[0].arn + : "arn:${data.aws_partition.current.partition}:ecr:${var.region}:${local.account_id}:repository/${var.ecr_repository_name}" + ) + runner_repository_url = ( + var.private_ecr + ? ( + var.create_ecr_repository + ? aws_ecr_repository.runner[0].repository_url + : "${local.account_id}.dkr.ecr.${var.region}.${data.aws_partition.current.dns_suffix}/${var.ecr_repository_name}" + ) + : null + ) + microvm_image_arn = "arn:${data.aws_partition.current.partition}:lambda:${var.region}:${local.account_id}:microvm-image:${var.image_name}" + + base_tags = merge(var.tags, { + "app" = "codeapi" + "component" = "lambda-microvm" + }) +} + +# -------------------------------------------------------------------------- +# ECR: arm64 runner image consumed by the Lambda MicroVM image builder +# -------------------------------------------------------------------------- +resource "aws_ecr_repository" "runner" { + count = var.private_ecr && var.create_ecr_repository ? 1 : 0 + name = var.ecr_repository_name + image_tag_mutability = "IMMUTABLE" + force_delete = var.ecr_force_delete + tags = local.base_tags + + encryption_configuration { + encryption_type = "AES256" + } + + image_scanning_configuration { + scan_on_push = true + } +} + +resource "aws_ecr_lifecycle_policy" "runner" { + count = var.private_ecr && var.create_ecr_repository ? 1 : 0 + repository = aws_ecr_repository.runner[0].name + policy = jsonencode({ + rules = [{ + rulePriority = 1 + description = "Retain the newest ${var.ecr_max_image_count} runner images" + selection = { + tagStatus = "any" + countType = "imageCountMoreThan" + countNumber = var.ecr_max_image_count + } + action = { + type = "expire" + } + }] + }) +} + +# -------------------------------------------------------------------------- +# S3: code-artifact bucket (the zip that create-microvm-image reads) +# -------------------------------------------------------------------------- +resource "aws_s3_bucket" "artifact" { + count = var.create_artifact_bucket ? 1 : 0 + # Region in the name: S3 bucket names are globally unique, so applying this + # module in a second region with the same name_prefix would otherwise collide. + bucket = "${var.name_prefix}-artifacts-${var.region}-${local.account_id}" + force_destroy = var.artifact_force_destroy + tags = local.base_tags +} + +resource "aws_s3_bucket_public_access_block" "artifact" { + count = var.create_artifact_bucket ? 1 : 0 + bucket = aws_s3_bucket.artifact[0].id + block_public_acls = true + block_public_policy = true + ignore_public_acls = true + restrict_public_buckets = true +} + +resource "aws_s3_bucket_server_side_encryption_configuration" "artifact" { + count = var.create_artifact_bucket ? 1 : 0 + bucket = aws_s3_bucket.artifact[0].id + rule { + # SSE-S3 (not SSE-KMS): the Lambda build role reads the artifact with only + # s3:GetObject, so a KMS-encrypted bucket would AccessDenied without a + # kms:Decrypt grant. Artifacts are non-sensitive (the runner image zip); + # upgrade to aws:kms + a kms:Decrypt statement on the build role if you need + # a customer-managed key. + apply_server_side_encryption_by_default { + sse_algorithm = "AES256" + } + } +} + +resource "aws_s3_bucket_lifecycle_configuration" "artifact" { + count = var.create_artifact_bucket ? 1 : 0 + bucket = aws_s3_bucket.artifact[0].id + rule { + id = "expire-build-artifacts" + status = "Enabled" + filter {} + expiration { + days = var.artifact_retention_days + } + abort_incomplete_multipart_upload { + days_after_initiation = 7 + } + } +} + +# -------------------------------------------------------------------------- +# S3: session-workspace checkpoint bucket +# The CodeAPI control plane (not the MicroVM) reads/writes these. Encrypted, +# versioned for forensic history, and lifecycle-expired since checkpoints are a +# resumable cache rather than a system of record. +# -------------------------------------------------------------------------- +resource "aws_s3_bucket" "checkpoint" { + bucket = "${var.name_prefix}-checkpoints-${var.region}-${local.account_id}" + force_destroy = var.checkpoint_force_destroy + tags = local.base_tags +} + +resource "aws_s3_bucket_public_access_block" "checkpoint" { + bucket = aws_s3_bucket.checkpoint.id + block_public_acls = true + block_public_policy = true + ignore_public_acls = true + restrict_public_buckets = true +} + +resource "aws_s3_bucket_versioning" "checkpoint" { + bucket = aws_s3_bucket.checkpoint.id + versioning_configuration { + status = "Enabled" + } +} + +resource "aws_s3_bucket_server_side_encryption_configuration" "checkpoint" { + bucket = aws_s3_bucket.checkpoint.id + rule { + # SSE-S3 so the checkpoint access policy needs no kms:Decrypt / + # kms:GenerateDataKey grant (the MinIO-compatible client reads/writes with + # plain S3 perms). Still encrypted at rest with AWS-managed keys; switch to + # aws:kms + KMS grants on checkpoint_access if you require a CMK. + apply_server_side_encryption_by_default { + sse_algorithm = "AES256" + } + } +} + +resource "aws_s3_bucket_lifecycle_configuration" "checkpoint" { + bucket = aws_s3_bucket.checkpoint.id + depends_on = [ + aws_s3_bucket_versioning.checkpoint, + ] + rule { + id = "expire-checkpoints" + status = "Enabled" + filter {} + expiration { + days = var.checkpoint_retention_days + } + noncurrent_version_expiration { + noncurrent_days = var.checkpoint_noncurrent_retention_days + } + abort_incomplete_multipart_upload { + days_after_initiation = 7 + } + } +} + +# -------------------------------------------------------------------------- +# CloudWatch Logs +# Build logs land at the EXACT path `/aws/lambda-microvms/` (hyphen). +# DO NOT "correct" this to the AWS docs' `/aws/lambda/microvms/...` (slash) — the +# docs are wrong. Verified empirically against a live account: every build's log +# group is the hyphen form and the slash path does not exist. Switching to it +# would stop matching the group the builder writes to and lose the build logs +# (which then surface as a `CREATE_FAILED` with an empty `stateReason` — +# undebuggable). Pre-creating it sets retention; Lambda also auto-creates it if +# absent. Runtime VM stdout needs BOTH a cloudWatch logging config on RunMicrovm +# AND an execution role, or it goes nowhere. +# -------------------------------------------------------------------------- +resource "aws_cloudwatch_log_group" "build" { + name = "/aws/lambda-microvms/${var.image_name}" + retention_in_days = var.log_retention_days + tags = local.base_tags +} + +resource "aws_cloudwatch_log_group" "runtime" { + name = "/${var.name_prefix}/runtime" + retention_in_days = var.log_retention_days + tags = local.base_tags +} + +# -------------------------------------------------------------------------- +# IAM: build role (assumed by Lambda during create/update-microvm-image) +# Trust MUST include sts:TagSession, and perms MUST include writes to the exact +# pre-created build log group plus s3:GetObject or the build FAILS with an empty +# stateReason (undebuggable). +# -------------------------------------------------------------------------- +data "aws_iam_policy_document" "build_trust" { + statement { + effect = "Allow" + actions = ["sts:AssumeRole", "sts:TagSession"] + principals { + type = "Service" + identifiers = ["lambda.amazonaws.com"] + } + } +} + +resource "aws_iam_role" "build" { + name = "${var.name_prefix}-build" + assume_role_policy = data.aws_iam_policy_document.build_trust.json + tags = local.base_tags +} + +data "aws_iam_policy_document" "build_perms" { + statement { + sid = "ArtifactRead" + effect = "Allow" + actions = ["s3:GetObject"] + resources = ["arn:${data.aws_partition.current.partition}:s3:::${local.artifact_bucket}/*"] + } + + statement { + sid = "BuildLogs" + effect = "Allow" + actions = ["logs:CreateLogStream", "logs:PutLogEvents"] + # The group is pre-created above. Stream actions need the `:*` suffix; `/*` + # does not match log-stream ARNs. + resources = ["${aws_cloudwatch_log_group.build.arn}:*"] + } + + dynamic "statement" { + for_each = var.private_ecr ? [1] : [] + content { + sid = "PrivateEcrAuth" + effect = "Allow" + actions = ["ecr:GetAuthorizationToken"] + resources = ["*"] + } + } + + dynamic "statement" { + for_each = var.private_ecr ? [1] : [] + content { + sid = "PrivateEcrRepositoryPull" + effect = "Allow" + actions = [ + "ecr:BatchCheckLayerAvailability", + "ecr:GetDownloadUrlForLayer", + "ecr:BatchGetImage", + ] + resources = [local.runner_repository_arn] + } + } +} + +resource "aws_iam_role_policy" "build" { + name = "build" + role = aws_iam_role.build.id + policy = data.aws_iam_policy_document.build_perms.json +} + +# -------------------------------------------------------------------------- +# IAM: execution role (RunMicrovm executionRoleArn) — logging-only. +# The MicroVM never needs S3/network creds: checkpoints flow through the control +# plane over the authed proxy, so keep this role least-privilege. +# -------------------------------------------------------------------------- +data "aws_iam_policy_document" "exec_trust" { + statement { + effect = "Allow" + actions = ["sts:AssumeRole", "sts:TagSession"] + principals { + type = "Service" + identifiers = ["lambda.amazonaws.com"] + } + } +} + +resource "aws_iam_role" "execution" { + name = "${var.name_prefix}-exec" + assume_role_policy = data.aws_iam_policy_document.exec_trust.json + tags = local.base_tags +} + +data "aws_iam_policy_document" "exec_perms" { + statement { + sid = "RuntimeLogs" + effect = "Allow" + actions = ["logs:CreateLogStream", "logs:PutLogEvents"] + # The group is Terraform-managed; the runtime may write only its streams. + resources = ["${aws_cloudwatch_log_group.runtime.arn}:*"] + } +} + +resource "aws_iam_role_policy" "execution" { + name = "runtime-logs" + role = aws_iam_role.execution.id + policy = data.aws_iam_policy_document.exec_perms.json +} + +# -------------------------------------------------------------------------- +# IAM: CodeAPI worker caller policy +# RunMicrovm has two dependent permissions that are easy to miss: +# iam:PassRole for the runtime execution role and lambda:PassNetworkConnector +# for every ingress/egress connector supplied on the request. +# -------------------------------------------------------------------------- +data "aws_iam_policy_document" "worker_microvm_control" { + statement { + sid = "OperateCodeapiMicrovms" + effect = "Allow" + actions = [ + "lambda:RunMicrovm", + "lambda:GetMicrovm", + "lambda:CreateMicrovmAuthToken", + "lambda:TerminateMicrovm", + ] + resources = [local.microvm_image_arn] + } + + statement { + sid = "PassMicrovmExecutionRole" + effect = "Allow" + actions = ["iam:PassRole"] + resources = [aws_iam_role.execution.arn] + condition { + test = "StringEquals" + variable = "iam:PassedToService" + values = ["lambda.amazonaws.com"] + } + } + + statement { + sid = "PassMicrovmNetworkConnectors" + effect = "Allow" + actions = ["lambda:PassNetworkConnector"] + # This permission-only action currently exposes no resource type. + resources = ["*"] + } +} + +resource "aws_iam_policy" "worker_microvm_control" { + name = "${var.name_prefix}-worker-microvm-control" + policy = data.aws_iam_policy_document.worker_microvm_control.json + tags = local.base_tags +} + +resource "aws_iam_role_policy_attachment" "worker_microvm_control" { + count = var.codeapi_worker_role_name == "" ? 0 : 1 + role = var.codeapi_worker_role_name + policy_arn = aws_iam_policy.worker_microvm_control.arn +} + +# -------------------------------------------------------------------------- +# IAM policy document for the CodeAPI control plane's checkpoint access. +# Attach to your CodeAPI task role (preferred) or the optional user below. +# -------------------------------------------------------------------------- +data "aws_iam_policy_document" "checkpoint_access" { + statement { + sid = "CheckpointObjects" + effect = "Allow" + actions = ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"] + resources = ["${aws_s3_bucket.checkpoint.arn}/*"] + } + statement { + sid = "CheckpointList" + effect = "Allow" + actions = ["s3:ListBucket"] + resources = [aws_s3_bucket.checkpoint.arn] + } +} + +resource "aws_iam_policy" "checkpoint_access" { + name = "${var.name_prefix}-checkpoint-access" + policy = data.aws_iam_policy_document.checkpoint_access.json + tags = local.base_tags +} + +resource "aws_iam_role_policy_attachment" "worker_checkpoint_access" { + count = var.codeapi_worker_role_name == "" ? 0 : 1 + role = var.codeapi_worker_role_name + policy_arn = aws_iam_policy.checkpoint_access.arn +} + +resource "aws_iam_user" "checkpoint" { + count = var.create_checkpoint_access_user ? 1 : 0 + name = "${var.name_prefix}-checkpoint" + tags = local.base_tags +} + +resource "aws_iam_user_policy_attachment" "checkpoint" { + count = var.create_checkpoint_access_user ? 1 : 0 + user = aws_iam_user.checkpoint[0].name + policy_arn = aws_iam_policy.checkpoint_access.arn +} + +resource "aws_iam_access_key" "checkpoint" { + count = var.create_checkpoint_access_user ? 1 : 0 + user = aws_iam_user.checkpoint[0].name +} diff --git a/docs/lambda-microvm/terraform/outputs.tf b/docs/lambda-microvm/terraform/outputs.tf new file mode 100644 index 0000000..f2f5923 --- /dev/null +++ b/docs/lambda-microvm/terraform/outputs.tf @@ -0,0 +1,62 @@ +output "region" { + value = var.region +} + +output "artifact_bucket" { + description = "S3 bucket for the code-artifact zip (feed to build-lambda-microvm-artifact.sh S3_URI)." + value = local.artifact_bucket +} + +output "runner_ecr_repository_url" { + description = "Set as ECR_URI for scripts/build-lambda-microvm-artifact.sh. Null when private_ecr is false." + value = local.runner_repository_url +} + +output "checkpoint_bucket" { + description = "S3 bucket for session checkpoints (CODEAPI_CHECKPOINT_BUCKET)." + value = aws_s3_bucket.checkpoint.id +} + +output "build_role_arn" { + description = "Pass to create-microvm-image.ts as --build-role." + value = aws_iam_role.build.arn +} + +output "execution_role_arn" { + description = "Set as LAMBDA_MICROVM_EXECUTION_ROLE_ARN so runtime VM stdout reaches CloudWatch." + value = aws_iam_role.execution.arn +} + +output "build_log_group" { + value = aws_cloudwatch_log_group.build.name +} + +output "runtime_log_group" { + value = aws_cloudwatch_log_group.runtime.name +} + +output "checkpoint_access_policy_arn" { + description = "Attach to the CodeAPI worker/task role for checkpoint S3 access." + value = aws_iam_policy.checkpoint_access.arn +} + +output "worker_microvm_control_policy_arn" { + description = "Attach to the CodeAPI worker/task role for Run/Get/CreateToken/Terminate plus PassRole/PassNetworkConnector." + value = aws_iam_policy.worker_microvm_control.arn +} + +output "microvm_image_arn" { + description = "Expected image ARN for image_name; create-microvm-image.ts creates this resource." + value = local.microvm_image_arn +} + +output "checkpoint_access_key_id" { + description = "Only when create_checkpoint_access_user = true. Use as MINIO_ACCESS_KEY." + value = try(aws_iam_access_key.checkpoint[0].id, null) +} + +output "checkpoint_secret_access_key" { + description = "Only when create_checkpoint_access_user = true. Use as MINIO_SECRET_KEY." + value = try(aws_iam_access_key.checkpoint[0].secret, null) + sensitive = true +} diff --git a/docs/lambda-microvm/terraform/terraform.tfvars.example b/docs/lambda-microvm/terraform/terraform.tfvars.example new file mode 100644 index 0000000..c978407 --- /dev/null +++ b/docs/lambda-microvm/terraform/terraform.tfvars.example @@ -0,0 +1,46 @@ +# Copy to terraform.tfvars and edit. + +region = "us-east-1" +name_prefix = "codeapi-microvm" + +# Must match the --name you pass to create-microvm-image (pre-creates the build +# log group at /aws/lambda-microvms/). +image_name = "codeapi-session" + +# Runner container image base. true = build role can pull from a private ECR repo +# in this account (the default runner Dockerfile does). false for public bases. +private_ecr = true + +# The module creates this repo by default. Set create_ecr_repository=false only +# when the repository already exists. +create_ecr_repository = true +ecr_repository_name = "codeapi-microvm-runner" +ecr_max_image_count = 30 + +# AIML-dev is a disposable integration stack, so this example explicitly opts +# into destructive teardown. Terraform defaults all three values to false. +# Keep them false for retained environments: `terraform destroy` will then fail +# instead of silently deleting non-empty buckets or a non-empty image repo. +artifact_force_destroy = true +checkpoint_force_destroy = true +ecr_force_delete = true + +# Optional: attach both emitted runtime policies directly to an existing +# CodeAPI worker/task role. Leave empty to attach the output ARNs elsewhere. +codeapi_worker_role_name = "" + +# Checkpoints are a resumable cache; keep current versions for 14 days and an +# S3 noncurrent version for one additional day after replacement/expiration. +artifact_retention_days = 30 +checkpoint_retention_days = 14 +checkpoint_noncurrent_retention_days = 1 +log_retention_days = 30 + +# Prefer attaching checkpoint_access_policy_arn to your CodeAPI task role. +# Only set true if you must hand static keys to a non-role deployment. +create_checkpoint_access_user = false + +tags = { + team = "code-interpreter" + env = "dev" +} diff --git a/docs/lambda-microvm/terraform/variables.tf b/docs/lambda-microvm/terraform/variables.tf new file mode 100644 index 0000000..caae4ce --- /dev/null +++ b/docs/lambda-microvm/terraform/variables.tf @@ -0,0 +1,171 @@ +variable "region" { + description = "AWS region to provision into. Lambda MicroVMs must be available here." + type = string + default = "us-east-1" +} + +variable "name_prefix" { + description = "Prefix for all created resource names (roles, buckets, log groups)." + type = string + default = "codeapi-microvm" +} + +variable "image_name" { + description = <<-EOT + Name of the MicroVM image you will create with the SDK/CLI helper. Only used + to pre-create the build log group at the exact path Lambda writes to + (`/aws/lambda-microvms/`). Must match the `--name` you pass to + create-microvm-image. + EOT + type = string + default = "codeapi-session" +} + +variable "create_artifact_bucket" { + description = <<-EOT + Create an S3 bucket to hold the code-artifact zip that + `scripts/build-lambda-microvm-artifact.sh` uploads. Set false to reuse an + existing bucket via `artifact_bucket_name`. + EOT + type = bool + default = true +} + +variable "artifact_bucket_name" { + description = "Existing artifact bucket name when create_artifact_bucket = false." + type = string + default = "" + + # Reject an empty name when reusing an existing bucket, else the build-role + # policy resolves to `arn:aws:s3:::/*` and the build can't read the artifact. + validation { + condition = var.create_artifact_bucket || length(var.artifact_bucket_name) > 0 + error_message = "artifact_bucket_name must be set when create_artifact_bucket is false." + } +} + +variable "artifact_retention_days" { + description = "Days to retain uploaded MicroVM build artifacts in the managed artifact bucket." + type = number + default = 30 + + validation { + condition = var.artifact_retention_days > 0 && floor(var.artifact_retention_days) == var.artifact_retention_days + error_message = "artifact_retention_days must be a positive whole number." + } +} + +variable "artifact_force_destroy" { + description = "Allow terraform destroy to remove a non-empty managed artifact bucket. Opt in only for disposable dev stacks." + type = bool + default = false +} + +variable "checkpoint_retention_days" { + description = <<-EOT + Days to keep session-workspace checkpoints in the checkpoint bucket before + lifecycle expiration. Checkpoints are a resumable cache, not a system of + record, so a short window is fine. + EOT + type = number + default = 14 + + validation { + condition = var.checkpoint_retention_days > 0 && floor(var.checkpoint_retention_days) == var.checkpoint_retention_days + error_message = "checkpoint_retention_days must be a positive whole number." + } +} + +variable "checkpoint_noncurrent_retention_days" { + description = <<-EOT + Days to keep a noncurrent checkpoint object version after S3 replaces or + expires its current version. This is separate from + checkpoint_retention_days so bucket versioning does not silently double the + intended resumable-cache retention window. + EOT + type = number + default = 1 + + validation { + condition = var.checkpoint_noncurrent_retention_days > 0 && floor(var.checkpoint_noncurrent_retention_days) == var.checkpoint_noncurrent_retention_days + error_message = "checkpoint_noncurrent_retention_days must be a positive whole number." + } +} + +variable "checkpoint_force_destroy" { + description = "Allow terraform destroy to remove a non-empty checkpoint bucket. Opt in only for disposable dev stacks." + type = bool + default = false +} + +variable "log_retention_days" { + description = "CloudWatch Logs retention for the build + runtime log groups." + type = number + default = 30 +} + +variable "private_ecr" { + description = <<-EOT + Grant the build role permission to pull the runner container image from a + private ECR repo in this account. Required when the code-artifact Dockerfile + uses `FROM .dkr.ecr...`. Leave false for public base images. + EOT + type = bool + default = true +} + +variable "create_ecr_repository" { + description = "Create the private ECR repository used by the runner image build." + type = bool + default = true +} + +variable "ecr_repository_name" { + description = "Private ECR repository containing the arm64 CodeAPI runner image." + type = string + default = "codeapi-microvm-runner" +} + +variable "ecr_force_delete" { + description = "Allow terraform destroy to remove a non-empty runner ECR repository. Opt in only for disposable dev stacks." + type = bool + default = false +} + +variable "ecr_max_image_count" { + description = "Maximum number of runner images retained by the managed ECR repository." + type = number + default = 30 + + validation { + condition = var.ecr_max_image_count > 0 && floor(var.ecr_max_image_count) == var.ecr_max_image_count + error_message = "ecr_max_image_count must be a positive whole number." + } +} + +variable "codeapi_worker_role_name" { + description = <<-EOT + Existing CodeAPI worker/task IAM role name. When set, Terraform attaches + both the MicroVM control-plane policy and checkpoint S3 policy. Leave empty + to consume the two policy ARN outputs in your deployment stack. + EOT + type = string + default = "" +} + +variable "create_checkpoint_access_user" { + description = <<-EOT + Create an IAM user + access key with read/write on the checkpoint bucket, for + the CodeAPI service's MinIO-compatible checkpoint client (MINIO_ACCESS_KEY / + MINIO_SECRET_KEY). Prefer an ECS task role, EC2 instance profile, or IRSA; + create static credentials only for a deployment that cannot use a role. + EOT + type = bool + default = false +} + +variable "tags" { + description = "Tags applied to every created resource." + type = map(string) + default = {} +} diff --git a/docs/lambda-microvm/terraform/versions.tf b/docs/lambda-microvm/terraform/versions.tf new file mode 100644 index 0000000..3d97e67 --- /dev/null +++ b/docs/lambda-microvm/terraform/versions.tf @@ -0,0 +1,15 @@ +terraform { + # >= 1.9 for cross-variable references in variable validation blocks. + required_version = ">= 1.9" + + required_providers { + aws = { + source = "hashicorp/aws" + version = ">= 5.40, < 7.0" + } + } +} + +provider "aws" { + region = var.region +} diff --git a/scripts/build-lambda-microvm-artifact.sh b/scripts/build-lambda-microvm-artifact.sh new file mode 100755 index 0000000..8e8869b --- /dev/null +++ b/scripts/build-lambda-microvm-artifact.sh @@ -0,0 +1,192 @@ +#!/bin/bash +# Builds the AWS Lambda MicroVM sandbox-runner artifacts. +# +# The Lambda MicroVM image pipeline is: zip(Dockerfile) -> S3 -> +# CreateMicrovmImage builds it on the AL2023 MicroVM base image. Our +# Dockerfile is a single FROM pointing at the prebuilt arm64 runner image +# in a same-account ECR repo (Lambda's build infra can pull it there). +# +# Stages (each optional, in order): +# build docker buildx the arm64 lambda-microvm-runner target (no AWS) +# push push to ECR (needs AWS_PROFILE + repo) +# zip render the code-artifact Dockerfile and zip it (no AWS) +# upload upload the zip to S3 (needs AWS_PROFILE + bucket) +# +# Usage: +# scripts/build-lambda-microvm-artifact.sh build +# scripts/build-lambda-microvm-artifact.sh build push zip upload +# +# Env: +# ECR_URI e.g. 951834775723.dkr.ecr.us-east-2.amazonaws.com/codeapi-microvm-runner +# IMAGE_TAG default: git short sha +# IMAGE_DIGEST sha256 digest for a separately-invoked zip stage (normally +# captured automatically by push) +# S3_URI e.g. s3://codeapi-microvm-artifacts/runner +# AWS_PROFILE e.g. librechat-dev +# AWS_REGION required for push/upload +set -euo pipefail + +cd "$(dirname "$0")/.." + +if [ -z "${IMAGE_TAG:-}" ]; then + if ! git diff --quiet || ! git diff --cached --quiet || [ -n "$(git ls-files --others --exclude-standard)" ]; then + echo "IMAGE_TAG is required when the worktree is dirty; use an explicit unique dev tag." >&2 + exit 1 + fi + IMAGE_TAG="$(git rev-parse --short HEAD)" +fi +ECR_URI="${ECR_URI:-}" +S3_URI="${S3_URI:-}" +OUT_DIR="${OUT_DIR:-.build-lambda-microvm}" +LOCAL_TAG="codeapi-lambda-microvm-runner:${IMAGE_TAG}" +IMAGE_DIGEST="${IMAGE_DIGEST:-}" + +require_ecr() { + [ -n "$ECR_URI" ] || { echo "ECR_URI is required for this stage" >&2; exit 1; } +} + +file_sha256() { + if command -v sha256sum >/dev/null 2>&1; then + sha256sum "$1" | awk '{print $1}' + else + shasum -a 256 "$1" | awk '{print $1}' + fi +} + +resolve_image_digest() { + if [ -z "$IMAGE_DIGEST" ] && [ -f "$OUT_DIR/image-digest" ]; then + local cached_repository cached_tag + cached_repository="$(sed -n '1p' "$OUT_DIR/image-repository" 2>/dev/null || true)" + cached_tag="$(sed -n '1p' "$OUT_DIR/image-tag" 2>/dev/null || true)" + if [ "$cached_repository" != "$ECR_URI" ] || [ "$cached_tag" != "$IMAGE_TAG" ]; then + echo "Cached image digest does not belong to ECR_URI=$ECR_URI IMAGE_TAG=$IMAGE_TAG; run push first or set IMAGE_DIGEST explicitly." >&2 + exit 1 + fi + IMAGE_DIGEST="$(sed -n '1p' "$OUT_DIR/image-digest")" + fi + local digest_hex="${IMAGE_DIGEST#sha256:}" + if [ "$digest_hex" = "$IMAGE_DIGEST" ] \ + || [ "${#digest_hex}" -ne 64 ] \ + || [[ "$digest_hex" == *[!0-9a-f]* ]]; then + echo "An immutable ECR IMAGE_DIGEST is required (run push first or set it explicitly)." >&2 + exit 1 + fi +} + +do_build() { + local tags=(-t "$LOCAL_TAG") + if [ -n "$ECR_URI" ]; then + tags+=(-t "$ECR_URI:$IMAGE_TAG") + fi + echo ">> buildx arm64 lambda-microvm-runner (${LOCAL_TAG})" + docker buildx build \ + --platform linux/arm64 \ + --target lambda-microvm-runner \ + -f api/Dockerfile \ + "${tags[@]}" \ + --load \ + . +} + +do_push() { + require_ecr + mkdir -p "$OUT_DIR" + echo ">> pushing $ECR_URI:$IMAGE_TAG" + aws ecr get-login-password --region "${AWS_REGION:?AWS_REGION required}" \ + | docker login --username AWS --password-stdin "${ECR_URI%%/*}" + # `build` is intentionally usable without AWS/ECR configuration. Tag here as + # well so a later, separately invoked `push` stage still has the remote tag. + docker image tag "$LOCAL_TAG" "$ECR_URI:$IMAGE_TAG" + docker push "$ECR_URI:$IMAGE_TAG" | tee "$OUT_DIR/push.log" + IMAGE_DIGEST="$(sed -n 's/^.*digest: \(sha256:[0-9a-f]\{64\}\).*$/\1/p' "$OUT_DIR/push.log" | tail -n 1)" + [ -n "$IMAGE_DIGEST" ] || { + echo "Could not determine the pushed ECR digest; refusing to render a mutable artifact." >&2 + exit 1 + } + printf '%s\n' "$IMAGE_DIGEST" > "$OUT_DIR/image-digest" + printf '%s\n' "$ECR_URI" > "$OUT_DIR/image-repository" + printf '%s\n' "$IMAGE_TAG" > "$OUT_DIR/image-tag" + echo ">> immutable runner ref: $ECR_URI@$IMAGE_DIGEST" +} + +do_zip() { + require_ecr + mkdir -p "$OUT_DIR" + resolve_image_digest + cat > "$OUT_DIR/Dockerfile" < "$OUT_DIR/artifact-image-repository" + printf '%s\n' "$IMAGE_TAG" > "$OUT_DIR/artifact-image-tag" + printf '%s\n' "$IMAGE_DIGEST" > "$OUT_DIR/artifact-image-digest" + file_sha256 "$OUT_DIR/artifact.zip" > "$OUT_DIR/artifact-sha256" + echo ">> wrote $OUT_DIR/artifact.zip (FROM ${ECR_URI}@${IMAGE_DIGEST})" +} + +do_upload() { + require_ecr + [ -n "$S3_URI" ] || { echo "S3_URI is required for upload" >&2; exit 1; } + resolve_image_digest + if ! { + [ -f "$OUT_DIR/artifact.zip" ] \ + && [ -f "$OUT_DIR/artifact-image-repository" ] \ + && [ -f "$OUT_DIR/artifact-image-tag" ] \ + && [ -f "$OUT_DIR/artifact-image-digest" ] \ + && [ -f "$OUT_DIR/artifact-sha256" ] + }; then + echo "No provenance-bound artifact found; run the zip stage first." >&2 + exit 1 + fi + local artifact_repository artifact_tag artifact_digest artifact_hash actual_hash + artifact_repository="$(sed -n '1p' "$OUT_DIR/artifact-image-repository")" + artifact_tag="$(sed -n '1p' "$OUT_DIR/artifact-image-tag")" + artifact_digest="$(sed -n '1p' "$OUT_DIR/artifact-image-digest")" + artifact_hash="$(sed -n '1p' "$OUT_DIR/artifact-sha256")" + actual_hash="$(file_sha256 "$OUT_DIR/artifact.zip")" + if [ "$artifact_repository" != "$ECR_URI" ] \ + || [ "$artifact_tag" != "$IMAGE_TAG" ] \ + || [ "$artifact_digest" != "$IMAGE_DIGEST" ] \ + || [ "$artifact_hash" != "$actual_hash" ]; then + echo "artifact.zip provenance does not match the current repository, tag, digest, or bytes; run zip again before upload." >&2 + exit 1 + fi + local key="$S3_URI/runner-${IMAGE_TAG}.zip" + aws s3 cp "$OUT_DIR/artifact.zip" "$key" --region "${AWS_REGION:?AWS_REGION required}" + echo ">> uploaded $key" + cat < \\ + --region \${AWS_REGION} \\ + --env-json "\$MICROVM_IMAGE_ENV_JSON" + +Notes: +- env vars (SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY, SANDBOX_REQUIRE_EGRESS_MANIFEST, + EGRESS_GATEWAY_URL, SANDBOX_ALLOWED_LOCAL_NETWORK_PORT, SANDBOX_FORWARD_TARGET) + are image-build-time config: pass them through MICROVM_IMAGE_ENV_JSON or + --env-json on the helper. +- Build the image HOOKLESS (no --hooks). Lambda's image build hooks only route + on the snapshot-compatible Lambda base container image, and enabling any + runtime hook forces the /ready build hook, which never reaches a stock + container's listener (builds then fail at the ready timeout). Session mode is + delivered per-request via the X-Runtime-Session-Id header instead; idle + suspend/resume is handled by RunMicrovm's native idlePolicy. +EOF +} + +for stage in "$@"; do + case "$stage" in + build) do_build ;; + push) do_push ;; + zip) do_zip ;; + upload) do_upload ;; + *) echo "Unknown stage: $stage (expected build|push|zip|upload)" >&2; exit 1 ;; + esac +done + +[ $# -gt 0 ] || { echo "No stages given. Usage: $0 build [push] [zip] [upload]" >&2; exit 1; } diff --git a/service/bun.lock b/service/bun.lock index e4d0f3c..ab7a9e1 100644 --- a/service/bun.lock +++ b/service/bun.lock @@ -5,6 +5,8 @@ "": { "name": "code-execution-lambda", "dependencies": { + "@aws-sdk/client-lambda-microvms": "3.1079.0", + "@aws-sdk/client-s3": "3.1079.0", "@opentelemetry/api": "1.9.1", "@opentelemetry/core": "2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "0.218.0", @@ -50,6 +52,44 @@ }, }, "packages": { + "@aws-sdk/checksums": ["@aws-sdk/checksums@3.1000.19", "", { "dependencies": { "@aws-sdk/core": "^3.976.0", "@aws-sdk/types": "^3.974.2", "@smithy/core": "^3.29.4", "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-Hc4N100RdkuWshKBnhPzmpdftfi9mCLz+OHFELHM1QIgMH4QRUUWyWgfiebta/YX2Bd62wTcm3EqAP8TeXv0gA=="], + + "@aws-sdk/client-lambda-microvms": ["@aws-sdk/client-lambda-microvms@3.1079.0", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/credential-provider-node": "^3.972.62", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/fetch-http-handler": "^5.6.2", "@smithy/node-http-handler": "^4.9.2", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-fLtosmk8h3cPasOob+vmOhRToVkNVcycsDi3x0ObKloNP1uNxDYWQVB3j7Z+xzfm8Wbk9Jq7A5Q/7e1rbAiNKw=="], + + "@aws-sdk/client-s3": ["@aws-sdk/client-s3@3.1079.0", "", { "dependencies": { "@aws-sdk/checksums": "^3.1000.12", "@aws-sdk/core": "^3.974.27", "@aws-sdk/credential-provider-node": "^3.972.62", "@aws-sdk/middleware-sdk-s3": "^3.972.58", "@aws-sdk/signature-v4-multi-region": "^3.996.38", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/fetch-http-handler": "^5.6.2", "@smithy/node-http-handler": "^4.9.2", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-di9U/7Po7qlVYb2dq58ULsbBAE1pBIk53rux+50LQCvH1X+/l1Ys+BIk/QLBtdaK1nADk0xRNEBbA1QWVnMccw=="], + + "@aws-sdk/core": ["@aws-sdk/core@3.974.27", "", { "dependencies": { "@aws-sdk/types": "^3.973.15", "@aws-sdk/xml-builder": "^3.972.33", "@aws/lambda-invoke-store": "^0.3.0", "@smithy/core": "^3.29.0", "@smithy/signature-v4": "^5.6.1", "@smithy/types": "^4.15.1", "bowser": "^2.11.0", "tslib": "^2.6.2" } }, "sha512-WRWEgIq6vx+NU6ot3VrRu4Jovj9MIObitSi6of/GV5THDDPccBhivCRNkWJutMM+m3GvdeI3l/UbGNcoOobxOA=="], + + "@aws-sdk/credential-provider-env": ["@aws-sdk/credential-provider-env@3.972.53", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-+KDA3uc/HZ1vIneGu5QMQb0gAXDYrm2vOE60+BJ7lS0YinMQ5i2oV4PR1A16XkF6K1IbSwjEHd1hQIIgMsK48w=="], + + "@aws-sdk/credential-provider-http": ["@aws-sdk/credential-provider-http@3.972.55", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/fetch-http-handler": "^5.6.2", "@smithy/node-http-handler": "^4.9.2", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-1gBfkWY3RWeBlCoB9lIJjXMx45/54wxcgfzv6BY9otTmMrZPcNPi1v+MwZxxaCUg441NV3jsr1efnFNCXiW70g=="], + + "@aws-sdk/credential-provider-ini": ["@aws-sdk/credential-provider-ini@3.972.60", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/credential-provider-env": "^3.972.53", "@aws-sdk/credential-provider-http": "^3.972.55", "@aws-sdk/credential-provider-login": "^3.972.59", "@aws-sdk/credential-provider-process": "^3.972.53", "@aws-sdk/credential-provider-sso": "^3.972.59", "@aws-sdk/credential-provider-web-identity": "^3.972.59", "@aws-sdk/nested-clients": "^3.997.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/credential-provider-imds": "^4.4.5", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-CV2md+PXvABwRjApWGhQ0wACy9WSFIhnUGrovLcjnjBCd/46TbuivLADtkF8IWNjtCQmQ+2IagSaxqBYqXBNAQ=="], + + "@aws-sdk/credential-provider-login": ["@aws-sdk/credential-provider-login@3.972.59", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/nested-clients": "^3.997.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-JG4S9yyA1GFzJdJXqLKrUzZbyK+VDp2QIsJD7YOicJHAhqymfHpDJIok2dLnhOdVB0I37RjdC53uOwCMVS00gw=="], + + "@aws-sdk/credential-provider-node": ["@aws-sdk/credential-provider-node@3.972.62", "", { "dependencies": { "@aws-sdk/credential-provider-env": "^3.972.53", "@aws-sdk/credential-provider-http": "^3.972.55", "@aws-sdk/credential-provider-ini": "^3.972.60", "@aws-sdk/credential-provider-process": "^3.972.53", "@aws-sdk/credential-provider-sso": "^3.972.59", "@aws-sdk/credential-provider-web-identity": "^3.972.59", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/credential-provider-imds": "^4.4.5", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-S6Slq3Tx7bvFk5yc34XNADyZYTX2HUXvaFAnowGRQnhjBO8J/mP62Fn7lxvJwjaDyYm/7gh9h6HEHaltRyMFXw=="], + + "@aws-sdk/credential-provider-process": ["@aws-sdk/credential-provider-process@3.972.53", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-EhfH+MQlqOMCkXIVa8MMObPzAQqwTTtxA7KhEJiyPeuNVA8PLOOUpgK7nBrgaDaGiIDLN/9LpGdaHuDjomeRTw=="], + + "@aws-sdk/credential-provider-sso": ["@aws-sdk/credential-provider-sso@3.972.59", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/nested-clients": "^3.997.27", "@aws-sdk/token-providers": "3.1079.0", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-h8793pOjcImx0SB+VcLONcaQQ57VAvKVuqyewQMRKqqH+CSXsG2dwOeLMUJPMxLdNvL7dXOM0ueTukyNUnu5mA=="], + + "@aws-sdk/credential-provider-web-identity": ["@aws-sdk/credential-provider-web-identity@3.972.59", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/nested-clients": "^3.997.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-VoyO9+vl3XVmpZwn4obskrWIkrA/Jf3lSe1E3ZERlaN9u0D4YZ6+HywC3+L98QOXqZesEfedk67gRER8tK8+8w=="], + + "@aws-sdk/middleware-sdk-s3": ["@aws-sdk/middleware-sdk-s3@3.972.65", "", { "dependencies": { "@aws-sdk/core": "^3.976.0", "@aws-sdk/signature-v4-multi-region": "^3.996.41", "@aws-sdk/types": "^3.974.2", "@smithy/core": "^3.29.4", "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-udwNhRfDTfCB98mAHjjgsnKQlxygB4e0X+Obne/XjJpvVsF0YCQC8ZErd/8Z6IPoLQjtiKHzwqEDbZiLrJEnOg=="], + + "@aws-sdk/nested-clients": ["@aws-sdk/nested-clients@3.997.27", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/signature-v4-multi-region": "^3.996.38", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/fetch-http-handler": "^5.6.2", "@smithy/node-http-handler": "^4.9.2", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-A8PIePF9NIIOJ/4Lg1rl9xm/+QaKkHGetq+Z9wb5B+3Da31YYXRo8n7IDMh5C+HQI5eyEmjrwkGWVdYtnLtbXQ=="], + + "@aws-sdk/signature-v4-multi-region": ["@aws-sdk/signature-v4-multi-region@3.996.38", "", { "dependencies": { "@aws-sdk/types": "^3.973.15", "@smithy/signature-v4": "^5.6.1", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-C379Sk+MiFZCfWZphKlMyLHKxV22OjoGM5KJjj5IJNJcOCWL4IGIpnEGzv1FQiRwhYXfq55SJMfxlqPE08JJ9g=="], + + "@aws-sdk/token-providers": ["@aws-sdk/token-providers@3.1079.0", "", { "dependencies": { "@aws-sdk/core": "^3.974.27", "@aws-sdk/nested-clients": "^3.997.27", "@aws-sdk/types": "^3.973.15", "@smithy/core": "^3.29.0", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-cbietrLlHPhhmbnMPTuDS4Zj/KNGhY+3vVhn6dwjO6Dqzrwothzg2srtcY34T9mlICsTXn34avDoWLHSntP54A=="], + + "@aws-sdk/types": ["@aws-sdk/types@3.973.15", "", { "dependencies": { "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-IULn8uBV/SMtmOIANsm4WHXIOtVPBWfOWs3WGL0j/sI+KhaYehvOw0ET+9urnn8MBpiijuU/0JOpuwKOE451PQ=="], + + "@aws-sdk/xml-builder": ["@aws-sdk/xml-builder@3.972.33", "", { "dependencies": { "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-ezbwz9WpuLctm6o7P2t2naDhVVPI5jFGrVefVybhcKGjU57VIyT46pQVO0RI2RYkUdhdj2Z9uSIlAzGZE9NW9A=="], + + "@aws/lambda-invoke-store": ["@aws/lambda-invoke-store@0.3.0", "", {}, "sha512-sl4Bm6yiMNYrZKkqqDFWN0UfnWhlS8ivKxrYl+6t0gCLrqr8y3B2IqZZbFRkfaVVp7C/baApyh71P+LeE1A2sQ=="], + "@colors/colors": ["@colors/colors@1.6.0", "", {}, "sha512-Ir+AOibqzrIsL6ajt3Rz3LskB7OiMVHqltZmspbW/TJuTVuyOMirVqAkjfY6JISiLHgyNqicAC8AyHHGzNd/dA=="], "@cspotcode/source-map-support": ["@cspotcode/source-map-support@0.8.1", "", { "dependencies": { "@jridgewell/trace-mapping": "0.3.9" } }, "sha512-IchNf6dN4tHoMFIn/7OE8LWZ19Y6q/67Bmf6vnGREv8RSbBVb9LPJxEcnwrcwX6ixSvaiGoomAUvu4YSxXrVgw=="], @@ -144,6 +184,18 @@ "@rtsao/scc": ["@rtsao/scc@1.1.0", "", {}, "sha512-zt6OdqaDoOnJ1ZYsCYGt9YmWzDXl4vQdKTyJev62gFhRGKdx7mcT54V9KIjg+d2wi9EXsPvAPKe7i7WjfVWB8g=="], + "@smithy/core": ["@smithy/core@3.29.1", "", { "dependencies": { "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-qoiY4nrk5OCu1+eIR1VB8l5DmON/oKiqrd5zZFAhXJXjJlLWQusKEW/SkBDAtGDcPaz86m9kfcE1lngU0GlM6A=="], + + "@smithy/credential-provider-imds": ["@smithy/credential-provider-imds@4.4.6", "", { "dependencies": { "@smithy/core": "^3.29.1", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-B2WQ/PV/H6Jeg3lrIq6bKUfa6Hy01mtK7CGs6lhjzHA6k4aagldH6T6eEjnzKl4HI0cJnAsxfJ19pgb5PV+CVQ=="], + + "@smithy/fetch-http-handler": ["@smithy/fetch-http-handler@5.6.3", "", { "dependencies": { "@smithy/core": "^3.29.1", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-CwCc/7SMTj45y97MUnDTbTaxvtAsiNNRm81z3abROIuMbMsC2Iy5EKfkkVdsKrz8WExQAAMx1EJapq+9j4fFTQ=="], + + "@smithy/node-http-handler": ["@smithy/node-http-handler@4.9.3", "", { "dependencies": { "@smithy/core": "^3.29.1", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-qZTa4gQFUo8RM02rk6q5UVTDLNrQ1oS20LsepBzqq1QBVc/EHJ03OOUADcqMZiXHArW+Y7+OGY0BpdTwZRq/Yg=="], + + "@smithy/signature-v4": ["@smithy/signature-v4@5.6.2", "", { "dependencies": { "@smithy/core": "^3.29.1", "@smithy/types": "^4.15.1", "tslib": "^2.6.2" } }, "sha512-QgHflghMoPxCJ9axiCVh8KZfbC9fuP6vkXXyK//E3cq7nLaSSyyLj0GAoqVWezYeDQmXIZhmlRvLE16jsqDK6g=="], + + "@smithy/types": ["@smithy/types@4.15.1", "", { "dependencies": { "tslib": "^2.6.2" } }, "sha512-x3L0XSACF6UYzKpa9biqiRMgvH5+wnFFew9Tm/grFYqgaupPwx/+ojDPpPJM8dZON3S9tjz5U+PQYsCBd1Mw5Q=="], + "@so-ric/colorspace": ["@so-ric/colorspace@1.1.6", "", { "dependencies": { "color": "^5.0.2", "text-hex": "1.0.x" } }, "sha512-/KiKkpHNOBgkFJwu9sh48LkHSMYGyuTcSFK/qMBdnOAlrRJzRSXAOFB5qwzaVQuDl8wAvHVMkaASQDReTahxuw=="], "@tsconfig/node10": ["@tsconfig/node10@1.0.12", "", {}, "sha512-UCYBaeFvM11aU2y3YPZ//O5Rhj+xKyzy7mvcIoAjASbigy8mHMryP5cK7dgjlz2hWxh1g5pLw084E0a/wlUSFQ=="], @@ -316,6 +368,8 @@ "body-parser": ["body-parser@1.20.4", "", { "dependencies": { "bytes": "~3.1.2", "content-type": "~1.0.5", "debug": "2.6.9", "depd": "2.0.0", "destroy": "~1.2.0", "http-errors": "~2.0.1", "iconv-lite": "~0.4.24", "on-finished": "~2.4.1", "qs": "~6.14.0", "raw-body": "~2.5.3", "type-is": "~1.6.18", "unpipe": "~1.0.0" } }, "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA=="], + "bowser": ["bowser@2.14.1", "", {}, "sha512-tzPjzCxygAKWFOJP011oxFHs57HzIhOEracIgAePE4pqB3LikALKnSzUyU4MGs9/iCEUuHlAJTjTc5M+u7YEGg=="], + "brace-expansion": ["brace-expansion@1.1.12", "", { "dependencies": { "balanced-match": "^1.0.0", "concat-map": "0.0.1" } }, "sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg=="], "braces": ["braces@3.0.3", "", { "dependencies": { "fill-range": "^7.1.1" } }, "sha512-yQbXgO/OSZVD2IsiLlro+7Hf6Q18EJrKSEsdoMzKePKXct3gvD8oLcOQdIzGupr5Fj+EDe8gO/lxc1BzfMpxvA=="], @@ -1014,6 +1068,24 @@ "yocto-queue": ["yocto-queue@0.1.0", "", {}, "sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q=="], + "@aws-sdk/checksums/@aws-sdk/core": ["@aws-sdk/core@3.976.0", "", { "dependencies": { "@aws-sdk/types": "^3.974.2", "@aws-sdk/xml-builder": "^3.972.36", "@aws/lambda-invoke-store": "^0.3.0", "@smithy/core": "^3.29.4", "@smithy/signature-v4": "^5.6.5", "@smithy/types": "^4.16.1", "bowser": "^2.11.0", "tslib": "^2.6.2" } }, "sha512-0cjRaEdlVoOrsNb9pP5q1Syyc8pXw5xSj2Np2ryReRTr9FppIIRVSdZK4lbnfmc2Hvgux/xBOUU6baB7z8//uA=="], + + "@aws-sdk/checksums/@aws-sdk/types": ["@aws-sdk/types@3.974.2", "", { "dependencies": { "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-3W6IUtSxFbH6X7Wb7DzGCV5QiFQsd0g8bOfntpmDxQlzBoKWUMBu/JPQR0DwkE+Hpnxd6db1tXbOwdeHddG6cA=="], + + "@aws-sdk/checksums/@smithy/core": ["@smithy/core@3.29.8", "", { "dependencies": { "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-rpCbCV+TimOBi3VLNBMmtTvgfOWcFIEAru3+TFlG87SL2F+te4jOnnNR+cf3uR4eJ5Qf4LnT80fqnBKgPRS6zA=="], + + "@aws-sdk/checksums/@smithy/types": ["@smithy/types@4.16.1", "", { "dependencies": { "tslib": "^2.6.2" } }, "sha512-0JFs3V2y2M9tKW5na/qxe69Zv+uxLMO7QBbhxF/FHu/Gp2NFZAAL9tWl9PU02xxo07pb3G9FTyjNc6D5uZrJIg=="], + + "@aws-sdk/middleware-sdk-s3/@aws-sdk/core": ["@aws-sdk/core@3.976.0", "", { "dependencies": { "@aws-sdk/types": "^3.974.2", "@aws-sdk/xml-builder": "^3.972.36", "@aws/lambda-invoke-store": "^0.3.0", "@smithy/core": "^3.29.4", "@smithy/signature-v4": "^5.6.5", "@smithy/types": "^4.16.1", "bowser": "^2.11.0", "tslib": "^2.6.2" } }, "sha512-0cjRaEdlVoOrsNb9pP5q1Syyc8pXw5xSj2Np2ryReRTr9FppIIRVSdZK4lbnfmc2Hvgux/xBOUU6baB7z8//uA=="], + + "@aws-sdk/middleware-sdk-s3/@aws-sdk/signature-v4-multi-region": ["@aws-sdk/signature-v4-multi-region@3.996.41", "", { "dependencies": { "@aws-sdk/types": "^3.974.2", "@smithy/signature-v4": "^5.6.5", "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-QMUytg+FQMGouc8gHS00KoYih3+N6cqmVI/pQGOIo7Nr7OpQaiXjSYOuL+vsPZ1tymY4LAQ8MYcHJmws5LRxng=="], + + "@aws-sdk/middleware-sdk-s3/@aws-sdk/types": ["@aws-sdk/types@3.974.2", "", { "dependencies": { "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-3W6IUtSxFbH6X7Wb7DzGCV5QiFQsd0g8bOfntpmDxQlzBoKWUMBu/JPQR0DwkE+Hpnxd6db1tXbOwdeHddG6cA=="], + + "@aws-sdk/middleware-sdk-s3/@smithy/core": ["@smithy/core@3.29.8", "", { "dependencies": { "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-rpCbCV+TimOBi3VLNBMmtTvgfOWcFIEAru3+TFlG87SL2F+te4jOnnNR+cf3uR4eJ5Qf4LnT80fqnBKgPRS6zA=="], + + "@aws-sdk/middleware-sdk-s3/@smithy/types": ["@smithy/types@4.16.1", "", { "dependencies": { "tslib": "^2.6.2" } }, "sha512-0JFs3V2y2M9tKW5na/qxe69Zv+uxLMO7QBbhxF/FHu/Gp2NFZAAL9tWl9PU02xxo07pb3G9FTyjNc6D5uZrJIg=="], + "@eslint/eslintrc/ignore": ["ignore@5.3.2", "", {}, "sha512-hsBTNUqQTDwkWtcdYI2i06Y/nUBEsNEDJKjWdigLvegy8kDuJAS8uRlpkkcQpyEXL0Z/pjDy5HBmMjRCJ2gq+g=="], "@types/minimatch/minimatch": ["minimatch@9.0.5", "", { "dependencies": { "brace-expansion": "^2.0.1" } }, "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow=="], @@ -1064,6 +1136,16 @@ "send/debug": ["debug@2.6.9", "", { "dependencies": { "ms": "2.0.0" } }, "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA=="], + "@aws-sdk/checksums/@aws-sdk/core/@aws-sdk/xml-builder": ["@aws-sdk/xml-builder@3.972.36", "", { "dependencies": { "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-RdGmS1GLrtaTOLE1ElSluMldNrpk9Emq6uYs8SS8iHlu5xTAmM9rRkM91o48+rIRryBtyO9t+uLYCoMG6jVMVA=="], + + "@aws-sdk/checksums/@aws-sdk/core/@smithy/signature-v4": ["@smithy/signature-v4@5.6.9", "", { "dependencies": { "@smithy/core": "^3.29.8", "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-g5rnEii/mkT0mjVJmlsaOfyNBtHNTecD9Lo4NP8D5HzMUEnZNpz7/FbvBCjNcV4vteHFAxOGiLUYNxPkDZZAPw=="], + + "@aws-sdk/middleware-sdk-s3/@aws-sdk/core/@aws-sdk/xml-builder": ["@aws-sdk/xml-builder@3.972.36", "", { "dependencies": { "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-RdGmS1GLrtaTOLE1ElSluMldNrpk9Emq6uYs8SS8iHlu5xTAmM9rRkM91o48+rIRryBtyO9t+uLYCoMG6jVMVA=="], + + "@aws-sdk/middleware-sdk-s3/@aws-sdk/core/@smithy/signature-v4": ["@smithy/signature-v4@5.6.9", "", { "dependencies": { "@smithy/core": "^3.29.8", "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-g5rnEii/mkT0mjVJmlsaOfyNBtHNTecD9Lo4NP8D5HzMUEnZNpz7/FbvBCjNcV4vteHFAxOGiLUYNxPkDZZAPw=="], + + "@aws-sdk/middleware-sdk-s3/@aws-sdk/signature-v4-multi-region/@smithy/signature-v4": ["@smithy/signature-v4@5.6.9", "", { "dependencies": { "@smithy/core": "^3.29.8", "@smithy/types": "^4.16.1", "tslib": "^2.6.2" } }, "sha512-g5rnEii/mkT0mjVJmlsaOfyNBtHNTecD9Lo4NP8D5HzMUEnZNpz7/FbvBCjNcV4vteHFAxOGiLUYNxPkDZZAPw=="], + "@types/minimatch/minimatch/brace-expansion": ["brace-expansion@2.0.2", "", { "dependencies": { "balanced-match": "^1.0.0" } }, "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ=="], "@typescript-eslint/typescript-estree/minimatch/brace-expansion": ["brace-expansion@2.0.2", "", { "dependencies": { "balanced-match": "^1.0.0" } }, "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ=="], diff --git a/service/package.json b/service/package.json index 89ca1ee..5a6b26d 100644 --- a/service/package.json +++ b/service/package.json @@ -26,6 +26,8 @@ "license": "Apache-2.0", "description": "", "dependencies": { + "@aws-sdk/client-lambda-microvms": "3.1079.0", + "@aws-sdk/client-s3": "3.1079.0", "@opentelemetry/api": "1.9.1", "@opentelemetry/core": "2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "0.218.0", diff --git a/service/scripts/create-microvm-image-lib.ts b/service/scripts/create-microvm-image-lib.ts new file mode 100644 index 0000000..542e1b6 --- /dev/null +++ b/service/scripts/create-microvm-image-lib.ts @@ -0,0 +1,190 @@ +import type { GetMicrovmImageOutput } from '@aws-sdk/client-lambda-microvms'; + +export type MicrovmImageSnapshot = Pick< + GetMicrovmImageOutput, + 'state' | 'imageArn' | 'latestActiveImageVersion' | 'latestFailedImageVersion' +>; + +export interface CompletedMicrovmImage { + state: 'CREATED' | 'UPDATED'; + imageArn: string; + imageVersion: string; + elapsedSeconds: number; +} + +export interface MicrovmImageListPage { + items?: Array<{ name?: string; imageArn?: string }>; + nextToken?: string; +} + +/** UpdateMicrovmImage requires a full ARN even though Create accepts a name. + * Resolve an exact name through every filtered page; never take the first + * substring match returned by `nameFilter`. */ +export async function resolveMicrovmImageArn( + nameOrArn: string, + listImages: ( + nameFilter: string, + nextToken?: string, + signal?: AbortSignal, + ) => Promise, + signal?: AbortSignal, +): Promise { + if (nameOrArn.startsWith('arn:')) return nameOrArn; + const matches = new Set(); + let nextToken: string | undefined; + do { + const page = await listImages(nameOrArn, nextToken, signal); + for (const image of page.items ?? []) { + if (image.name === nameOrArn && image.imageArn) matches.add(image.imageArn); + } + nextToken = page.nextToken; + } while (nextToken); + if (matches.size === 0) { + throw new Error(`MicroVM image "${nameOrArn}" was not found for update`); + } + if (matches.size > 1) { + throw new Error(`MicroVM image name "${nameOrArn}" resolved to multiple ARNs`); + } + return Array.from(matches)[0]; +} + +interface WaitForMicrovmImageOptions { + imageIdentifier: string; + deadlineMinutes: number; + getImage: (imageIdentifier: string, signal?: AbortSignal) => Promise; + now?: () => number; + sleep?: (milliseconds: number) => Promise; + pollIntervalMs?: number; + onPending?: (state: string, elapsedSeconds: number) => void; + /** Optional shared provisioning deadline. Passing the timestamp created + * before Create/Update makes those calls consume the same overall budget. */ + deadlineAtMs?: number; + startedAtMs?: number; + signal?: AbortSignal; +} + +export function positiveFiniteNumber(raw: string, label: string): number { + const value = Number(raw); + if (!Number.isFinite(value) || value <= 0) { + throw new Error(`${label} must be a positive number`); + } + return value; +} + +export function positiveInteger(raw: string, label: string): number { + const value = positiveFiniteNumber(raw, label); + if (!Number.isInteger(value)) { + throw new Error(`${label} must be a positive whole number`); + } + return value; +} + +export function parseStringMapJson(raw: string, label: string): Record { + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + throw new Error(`${label} must be valid JSON`); + } + if (parsed == null || Array.isArray(parsed) || typeof parsed !== 'object') { + throw new Error(`${label} must be a JSON object`); + } + const result: Record = {}; + for (const [key, value] of Object.entries(parsed)) { + if (key.length === 0 || typeof value !== 'string') { + throw new Error(`${label} must contain only non-empty keys with string values`); + } + result[key] = value; + } + return result; +} + +export async function waitForMicrovmImage( + options: WaitForMicrovmImageOptions, +): Promise { + if (!options.imageIdentifier.trim()) { + throw new Error('Create/UpdateMicrovmImage response omitted imageArn'); + } + if (!Number.isFinite(options.deadlineMinutes) || options.deadlineMinutes <= 0) { + throw new Error('MICROVM_BUILD_DEADLINE_MINUTES must be a positive number'); + } + + const now = options.now ?? Date.now; + const sleep = options.sleep ?? ((milliseconds: number) => new Promise( + (resolve) => setTimeout(resolve, milliseconds), + )); + const pollIntervalMs = options.pollIntervalMs ?? 20_000; + const startedAt = options.startedAtMs ?? now(); + const deadlineAt = options.deadlineAtMs + ?? startedAt + options.deadlineMinutes * 60_000; + let lastState = 'UNKNOWN'; + + for (;;) { + const beforePoll = now(); + if (beforePoll >= deadlineAt) { + const elapsedSeconds = Math.round((beforePoll - startedAt) / 1000); + throw new Error( + `Timed out after ${elapsedSeconds}s still in state ${lastState}. Check the build log group.`, + ); + } + + const pollController = new AbortController(); + const relayAbort = (): void => pollController.abort(options.signal?.reason); + if (options.signal?.aborted) relayAbort(); + else options.signal?.addEventListener('abort', relayAbort, { once: true }); + let timeout: ReturnType | undefined; + const remainingMs = Math.max(1, deadlineAt - beforePoll); + const deadline = new Promise((_resolve, reject) => { + timeout = setTimeout(() => { + const error = new Error( + `Timed out waiting for GetMicrovmImage in state ${lastState}. Check the build log group.`, + ); + pollController.abort(error); + reject(error); + }, remainingMs); + timeout.unref?.(); + }); + let image: MicrovmImageSnapshot; + try { + image = await Promise.race([ + options.getImage(options.imageIdentifier, pollController.signal), + deadline, + ]); + } finally { + if (timeout) clearTimeout(timeout); + options.signal?.removeEventListener('abort', relayAbort); + } + const currentTime = now(); + const elapsedSeconds = Math.round((currentTime - startedAt) / 1000); + const state = image.state ?? 'UNKNOWN'; + lastState = state; + + if (state === 'CREATED' || state === 'UPDATED') { + if (!image.imageArn || !image.latestActiveImageVersion) { + throw new Error(`${state} response omitted imageArn or latestActiveImageVersion`); + } + return { + state, + imageArn: image.imageArn, + imageVersion: image.latestActiveImageVersion, + elapsedSeconds, + }; + } + + if (state.includes('FAILED')) { + throw new Error( + `${state} after ${elapsedSeconds}s. failed version: ` + + `${image.latestFailedImageVersion || '(unknown — check the build log group)'}`, + ); + } + + if (currentTime >= deadlineAt) { + throw new Error( + `Timed out after ${elapsedSeconds}s still in state ${state}. Check the build log group.`, + ); + } + + options.onPending?.(state, elapsedSeconds); + await sleep(Math.min(pollIntervalMs, Math.max(0, deadlineAt - currentTime))); + } +} diff --git a/service/scripts/create-microvm-image.test.ts b/service/scripts/create-microvm-image.test.ts new file mode 100644 index 0000000..a292ffc --- /dev/null +++ b/service/scripts/create-microvm-image.test.ts @@ -0,0 +1,188 @@ +import { describe, expect, test } from 'bun:test'; +import { + positiveFiniteNumber, + positiveInteger, + parseStringMapJson, + resolveMicrovmImageArn, + waitForMicrovmImage, + type MicrovmImageSnapshot, +} from './create-microvm-image-lib'; + +function scriptedPoller(snapshots: MicrovmImageSnapshot[]): { + getImage: (imageIdentifier: string) => Promise; + identifiers: string[]; +} { + const queue = [...snapshots]; + const identifiers: string[] = []; + return { + identifiers, + getImage: async (imageIdentifier) => { + identifiers.push(imageIdentifier); + const next = queue.shift(); + if (!next) throw new Error('unexpected extra poll'); + return next; + }, + }; +} + +describe('waitForMicrovmImage', () => { + test('returns the pinned latest active version after create', async () => { + const poller = scriptedPoller([ + { state: 'CREATING', imageArn: 'arn:image' }, + { + state: 'CREATED', + imageArn: 'arn:image', + latestActiveImageVersion: '7', + // Guard against regressing to the nonexistent GetMicrovmImage + // `imageVersion` field that the first helper implementation read. + ...({ imageVersion: 'wrong-field' } as Record), + }, + ]); + let now = 0; + + const result = await waitForMicrovmImage({ + imageIdentifier: 'arn:image', + deadlineMinutes: 1, + getImage: poller.getImage, + now: () => now, + sleep: async (milliseconds) => { + now += milliseconds; + }, + pollIntervalMs: 20_000, + }); + + expect(result).toEqual({ + state: 'CREATED', + imageArn: 'arn:image', + imageVersion: '7', + elapsedSeconds: 20, + }); + expect(poller.identifiers).toEqual(['arn:image', 'arn:image']); + }); + + test('accepts the UPDATED terminal state', async () => { + const poller = scriptedPoller([ + { state: 'UPDATED', imageArn: 'arn:image', latestActiveImageVersion: '8' }, + ]); + + await expect(waitForMicrovmImage({ + imageIdentifier: 'arn:image', + deadlineMinutes: 1, + getImage: poller.getImage, + })).resolves.toMatchObject({ state: 'UPDATED', imageVersion: '8' }); + }); + + test('surfaces the failed image version', async () => { + const poller = scriptedPoller([ + { + state: 'UPDATE_FAILED', + imageArn: 'arn:image', + latestFailedImageVersion: '9', + }, + ]); + + await expect(waitForMicrovmImage({ + imageIdentifier: 'arn:image', + deadlineMinutes: 1, + getImage: poller.getImage, + })).rejects.toThrow('UPDATE_FAILED after 0s. failed version: 9'); + }); + + test('bounds a build that never leaves a pending state', async () => { + let now = 0; + + await expect(waitForMicrovmImage({ + imageIdentifier: 'arn:image', + deadlineMinutes: 1 / 60, + getImage: async () => ({ state: 'CREATING', imageArn: 'arn:image' }), + now: () => now, + sleep: async (milliseconds) => { + now += milliseconds; + }, + pollIntervalMs: 1_000, + })).rejects.toThrow('Timed out after 1s still in state CREATING'); + }); + + test('bounds a GetMicrovmImage call whose socket never settles', async () => { + const deadlineAtMs = Date.now() + 20; + let observedAbort = false; + await expect(waitForMicrovmImage({ + imageIdentifier: 'arn:image', + deadlineMinutes: 1, + deadlineAtMs, + getImage: async (_identifier, signal) => new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => { + observedAbort = true; + reject(signal.reason); + }, { once: true }); + }), + })).rejects.toThrow(); + expect(observedAbort).toBe(true); + }); + + test('rejects invalid numeric configuration instead of disabling the deadline', () => { + expect(() => positiveFiniteNumber('not-a-number', 'MICROVM_BUILD_DEADLINE_MINUTES')) + .toThrow('MICROVM_BUILD_DEADLINE_MINUTES must be a positive number'); + expect(() => positiveFiniteNumber('0', 'MICROVM_MEMORY_MIB')) + .toThrow('MICROVM_MEMORY_MIB must be a positive number'); + expect(() => positiveInteger('4096.5', 'MICROVM_MEMORY_MIB')) + .toThrow('MICROVM_MEMORY_MIB must be a positive whole number'); + }); + + test('accepts only a JSON object whose environment values are strings', () => { + expect(parseStringMapJson('{"PORT":"8080","FEATURE":"true"}', 'ENV')) + .toEqual({ PORT: '8080', FEATURE: 'true' }); + expect(() => parseStringMapJson('not-json', 'ENV')).toThrow('ENV must be valid JSON'); + expect(() => parseStringMapJson('[]', 'ENV')).toThrow('ENV must be a JSON object'); + expect(() => parseStringMapJson('{"PORT":8080}', 'ENV')) + .toThrow('ENV must contain only non-empty keys with string values'); + expect(() => parseStringMapJson('{"": "value"}', 'ENV')) + .toThrow('ENV must contain only non-empty keys with string values'); + }); +}); + +describe('resolveMicrovmImageArn', () => { + test('paginates and selects an exact name instead of a substring match', async () => { + const tokens: Array = []; + const arn = await resolveMicrovmImageArn( + 'codeapi-session', + async (_filter, nextToken) => { + tokens.push(nextToken); + if (!nextToken) { + return { + items: [{ + name: 'codeapi-session-old', + imageArn: 'arn:aws:lambda:us-east-1:1:microvm-image:old', + }], + nextToken: 'page-2', + }; + } + return { + items: [{ + name: 'codeapi-session', + imageArn: 'arn:aws:lambda:us-east-1:1:microvm-image:codeapi-session', + }], + }; + }, + ); + + expect(arn).toBe('arn:aws:lambda:us-east-1:1:microvm-image:codeapi-session'); + expect(tokens).toEqual([undefined, 'page-2']); + }); + + test('passes an ARN through without listing', async () => { + let listed = false; + const arn = 'arn:aws:lambda:us-east-1:1:microvm-image:codeapi-session'; + expect(await resolveMicrovmImageArn(arn, async () => { + listed = true; + return {}; + })).toBe(arn); + expect(listed).toBe(false); + }); + + test('fails closed when the filtered result has no exact match', async () => { + await expect(resolveMicrovmImageArn('wanted', async () => ({ + items: [{ name: 'wanted-old', imageArn: 'arn:old' }], + }))).rejects.toThrow('was not found for update'); + }); +}); diff --git a/service/scripts/create-microvm-image.ts b/service/scripts/create-microvm-image.ts new file mode 100644 index 0000000..c7ddeb7 --- /dev/null +++ b/service/scripts/create-microvm-image.ts @@ -0,0 +1,194 @@ +/** + * Create (or update) a hookless Lambda MicroVM image for the CodeAPI runner and + * poll it to CREATED. This is the one provisioning step Terraform can't own yet + * (the lambda-microvms service has no TF resource), so it lives here as a thin, + * guaranteed-correct wrapper over the SDK call proven during the spike. + * + * Hookless by design: Lambda's image build hooks only route on the + * snapshot-compatible Lambda base *container* image, and enabling any runtime + * hook forces the /ready build hook (which never reaches a stock container's + * listener, so the build fails at the ready timeout). Session mode is delivered + * per request via the X-Runtime-Session-Id header instead — no hooks needed. + * + * Run from the service workspace so the SDK resolves: + * cd service && AWS_PROFILE=... bun scripts/create-microvm-image.ts \ + * --name codeapi-session \ + * --artifact s3:///runner/runner-.zip \ + * --build-role arn:aws:iam:::role/codeapi-microvm-build \ + * --region us-east-1 + * + * Flags (or the UPPER_SNAKE env equivalents): + * --name MICROVM_IMAGE_NAME image name (default codeapi-session) + * --artifact MICROVM_ARTIFACT_URI s3:// uri of the code-artifact zip (required) + * --build-role MICROVM_BUILD_ROLE_ARN build role arn (required) + * --base-image MICROVM_BASE_IMAGE_ARN default arn:aws:lambda::aws:microvm-image:al2023-1 + * --base-version MICROVM_BASE_IMAGE_VERSION optional immutable managed-base version + * --region MICROVM_REGION default us-east-1 + * --memory MICROVM_MEMORY_MIB baseline memory (default 4096; RunMicrovm has no + * per-session memory override, so this image-time value + * is the ONLY memory lever — embedded-engine workloads + * like chdb OOM inside 2048) + * --update MICROVM_UPDATE=true update an existing image (new version) instead of create + */ +import { + LambdaMicrovmsClient, + CreateMicrovmImageCommand, + UpdateMicrovmImageCommand, + GetMicrovmImageCommand, + ListMicrovmImagesCommand, +} from '@aws-sdk/client-lambda-microvms'; +import { + positiveFiniteNumber, + positiveInteger, + parseStringMapJson, + resolveMicrovmImageArn, + waitForMicrovmImage, +} from './create-microvm-image-lib'; + +function arg(flag: string, env: string, fallback?: string): string | undefined { + const i = process.argv.indexOf(flag); + if (i >= 0 && process.argv[i + 1]) return process.argv[i + 1]; + return process.env[env] ?? fallback; +} + +const region = arg('--region', 'MICROVM_REGION', 'us-east-1') as string; +const name = arg('--name', 'MICROVM_IMAGE_NAME', 'codeapi-session') as string; +const artifactUri = arg('--artifact', 'MICROVM_ARTIFACT_URI'); +const buildRoleArn = arg('--build-role', 'MICROVM_BUILD_ROLE_ARN'); +const baseImageArn = arg( + '--base-image', + 'MICROVM_BASE_IMAGE_ARN', + `arn:aws:lambda:${region}:aws:microvm-image:al2023-1`, +) as string; +const baseImageVersion = arg('--base-version', 'MICROVM_BASE_IMAGE_VERSION'); +const memory = positiveInteger( + arg('--memory', 'MICROVM_MEMORY_MIB', '4096') as string, + 'MICROVM_MEMORY_MIB', +); +const buildDeadlineMinutes = positiveFiniteNumber( + process.env.MICROVM_BUILD_DEADLINE_MINUTES ?? '30', + 'MICROVM_BUILD_DEADLINE_MINUTES', +); +const isUpdate = (arg('--update', 'MICROVM_UPDATE') ?? 'false') === 'true' || process.argv.includes('--update'); + +if (!artifactUri || !buildRoleArn) { + console.error('Missing required --artifact and/or --build-role .'); + process.exit(2); +} + +/* Runner env is baked at image-build time (RunMicrovm does not inject it later), + * so the runner needs its egress-gateway / manifest config HERE or it can't + * upload outputs or verify execution manifests. By-reference inputs are pushed + * through the authenticated control plane, so FILE_SERVER_URL belongs only on + * the worker. The helper can't know your deployment's URLs or verifier key, so + * pass them via --env-json '{"EGRESS_GATEWAY_URL":"...", ...}' (or the + * MICROVM_IMAGE_ENV_JSON env). Typical keys: EGRESS_GATEWAY_URL, + * SANDBOX_ALLOWED_LOCAL_NETWORK_PORT, SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY, + * SANDBOX_REQUIRE_EGRESS_MANIFEST, CODEAPI_HARDENED_SANDBOX_MODE, and + * SANDBOX_FORWARD_TARGET. + * Runner limits also inject here and default LOW for session workloads: + * SANDBOX_RUN_TIMEOUT / SANDBOX_RUN_CPU_TIME (30000ms default kills long + * computations) and SANDBOX_OUTPUT_MAX_SIZE (1024 bytes default truncates + * output after a single verbose stack trace). */ +function parseEnvJson(): Record { + const raw = arg('--env-json', 'MICROVM_IMAGE_ENV_JSON'); + if (!raw) return {}; + return parseStringMapJson(raw, 'MICROVM_IMAGE_ENV_JSON'); +} + +/* Hard-won working image config (see docs/lambda-microvm/README.md): + * - additionalOsCapabilities ["ALL"]: nsjail needs CAP_SYS_ADMIN for its /proc + * mount inside the guest, else EPERM. + * - SANDBOX_USE_CGROUPV2=false: the app container can't read the cgroup v2 + * subtree_control, so fall back to rlimit-based caps. + * - NO hooks: hookless is the reliable build path (see header). */ +const shared = { + baseImageArn, + baseImageVersion, + buildRoleArn, + codeArtifact: { uri: artifactUri }, + cpuConfigurations: [{ architecture: 'ARM_64' as const }], + resources: [{ minimumMemoryInMiB: memory }], + additionalOsCapabilities: ['ALL' as const], + environmentVariables: { SANDBOX_USE_CGROUPV2: 'false', ...parseEnvJson() }, +}; + +const client = new LambdaMicrovmsClient({ region, retryMode: 'adaptive', maxAttempts: 3 }); + +async function main(): Promise { + console.log(`${isUpdate ? 'Updating' : 'Creating'} hookless MicroVM image "${name}" in ${region}...`); + console.log(` managed base: ${baseImageArn}${baseImageVersion ? ` @ ${baseImageVersion}` : ' (latest version)'}`); + const startedAtMs = Date.now(); + const deadlineAtMs = startedAtMs + buildDeadlineMinutes * 60_000; + const controller = new AbortController(); + const deadline = setTimeout(() => { + controller.abort(new Error( + `MicroVM image provisioning exceeded ${buildDeadlineMinutes} minute(s)`, + )); + }, Math.max(1, deadlineAtMs - Date.now())); + deadline.unref?.(); + let createdArn: string | undefined; + try { + if (isUpdate) { + const imageIdentifier = await resolveMicrovmImageArn( + name, + async (nameFilter, nextToken, signal) => client.send( + new ListMicrovmImagesCommand({ nameFilter, nextToken, maxResults: 50 }), + { abortSignal: signal }, + ), + controller.signal, + ); + const res = await client.send( + new UpdateMicrovmImageCommand({ imageIdentifier, ...shared }), + { abortSignal: controller.signal }, + ); + createdArn = res.imageArn ?? imageIdentifier; + } else { + const res = await client.send( + new CreateMicrovmImageCommand({ name, description: 'CodeAPI hookless session runner', ...shared }), + { abortSignal: controller.signal }, + ); + createdArn = res.imageArn; + } + + /* GetMicrovmImage and UpdateMicrovmImage require the full ARN. Create + * accepts a name; update resolves that name above. Create/Update, every + * lookup/poll, and every sleep share one hard deadline. */ + const completed = await waitForMicrovmImage({ + imageIdentifier: createdArn ?? '', + deadlineMinutes: buildDeadlineMinutes, + startedAtMs, + deadlineAtMs, + signal: controller.signal, + getImage: (imageIdentifier, signal) => client.send( + new GetMicrovmImageCommand({ imageIdentifier }), + { abortSignal: signal }, + ), + onPending: (state, elapsedSeconds) => { + if (elapsedSeconds % 60 < 20) console.log(` [${elapsedSeconds}s] ${state}`); + }, + }); + + console.log(`\n${completed.state} in ${completed.elapsedSeconds}s`); + console.log(` imageArn: ${completed.imageArn}`); + console.log(` version: ${completed.imageVersion}`); + console.log('\nSet on the CodeAPI service:'); + console.log(` LAMBDA_MICROVM_IMAGE_ARN=${completed.imageArn}`); + console.log(` LAMBDA_MICROVM_IMAGE_VERSION=${completed.imageVersion}`); + } catch (error) { + if (controller.signal.aborted) { + throw new Error( + `MicroVM image provisioning timed out after ${buildDeadlineMinutes} minute(s)`, + { cause: error }, + ); + } + throw error; + } finally { + clearTimeout(deadline); + } +} + +main().catch((error) => { + console.error('create-microvm-image failed:', error instanceof Error ? error.message : error); + process.exit(1); +}); diff --git a/service/src/config.spec.ts b/service/src/config.spec.ts index 665cff1..98056f5 100644 --- a/service/src/config.spec.ts +++ b/service/src/config.spec.ts @@ -1,5 +1,12 @@ import { describe, expect, it } from 'bun:test'; -import { languageConfig, resolveEgressGrantTtlSeconds, resolveLanguage } from './config'; +import { + jobDeadlineAtMs, + languageConfig, + lambdaMicrovmNumericConfigError, + resolveEgressGrantTtlSeconds, + resolveLanguage, + resolveLambdaMicrovmNumericConfig, +} from './config'; import { Languages } from './enum'; import { createPayload } from './payload'; import type { AuthenticatedRequest } from './types'; @@ -89,3 +96,70 @@ describe('egress grant TTL configuration', () => { expect(resolveEgressGrantTtlSeconds('not-a-number', 300000)).toBe(900); }); }); + +describe('job deadline accounting', () => { + it('counts time spent waiting in BullMQ against JOB_TIMEOUT', () => { + expect(jobDeadlineAtMs(1_000, 300_000, 50_000)).toBe(301_000); + }); + + it('falls back to worker start only when enqueue time is unavailable', () => { + expect(jobDeadlineAtMs(undefined, 300_000, 50_000)).toBe(350_000); + expect(jobDeadlineAtMs(Number.NaN, 300_000, 50_000)).toBe(350_000); + }); +}); + +describe('Lambda MicroVM numeric configuration', () => { + it('uses defaults only for absent or blank values and preserves suspend=0', () => { + expect(resolveLambdaMicrovmNumericConfig({})).toMatchObject({ + LAMBDA_MICROVM_PORT: 8080, + LAMBDA_MICROVM_MAX_DURATION_SECONDS: 28_800, + LAMBDA_MICROVM_IDLE_SECONDS: 1_800, + LAMBDA_MICROVM_SUSPEND_SECONDS: 1_800, + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: 300, + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: 60_000, + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: 5_000, + LAMBDA_MICROVM_LAUNCH_TPS: 4, + LAMBDA_MICROVM_TOKEN_TPS: 8, + }); + expect(resolveLambdaMicrovmNumericConfig({ + LAMBDA_MICROVM_SUSPEND_SECONDS: '0', + LAMBDA_MICROVM_PORT: ' ', + }).LAMBDA_MICROVM_SUSPEND_SECONDS).toBe(0); + }); + + it('accepts the supported boundary values', () => { + const numeric = resolveLambdaMicrovmNumericConfig({ + LAMBDA_MICROVM_PORT: '1', + LAMBDA_MICROVM_MAX_DURATION_SECONDS: '28800', + LAMBDA_MICROVM_IDLE_SECONDS: '60', + LAMBDA_MICROVM_SUSPEND_SECONDS: '0', + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: '1', + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: '1', + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: '1', + LAMBDA_MICROVM_LAUNCH_TPS: '1', + LAMBDA_MICROVM_TOKEN_TPS: '1', + }); + expect(lambdaMicrovmNumericConfigError(numeric)).toBeUndefined(); + }); + + it('rejects non-integers and values outside the supported ranges', () => { + const cases: Array<[string, string, string]> = [ + ['LAMBDA_MICROVM_PORT', '0', 'between 1 and 65535'], + ['LAMBDA_MICROVM_PORT', '65536', 'between 1 and 65535'], + ['LAMBDA_MICROVM_MAX_DURATION_SECONDS', '28801', 'between 1 and 28800'], + ['LAMBDA_MICROVM_IDLE_SECONDS', '59', 'between 60 and 28800'], + ['LAMBDA_MICROVM_SUSPEND_SECONDS', '-1', 'between 0 and 28800'], + ['LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS', '901', 'between 1 and 900'], + ['LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS', '0', 'at least 1'], + ['LAMBDA_MICROVM_HEALTH_TIMEOUT_MS', 'NaN', 'at least 1'], + ['LAMBDA_MICROVM_LAUNCH_TPS', '1.5', 'at least 1'], + ['LAMBDA_MICROVM_TOKEN_TPS', '-1', 'at least 1'], + ]; + + for (const [name, value, expected] of cases) { + const numeric = resolveLambdaMicrovmNumericConfig({ [name]: value }); + expect(lambdaMicrovmNumericConfigError(numeric)).toContain(name); + expect(lambdaMicrovmNumericConfigError(numeric)).toContain(expected); + } + }); +}); diff --git a/service/src/config.ts b/service/src/config.ts index 5abc57c..bfc7399 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -48,6 +48,131 @@ const defaultJobTimeoutMs = Number(process.env.JOB_TIMEOUT) || 300000; const defaultMaxFileSize = Number(process.env.MAX_FILE_SIZE) || 25 * 1024 * 1024; const defaultExecutionManifestTtlSeconds = Math.min(Math.ceil((defaultJobTimeoutMs + 60000) / 1000), 600); const EGRESS_GRANT_GRACE_MS = 10 * 60 * 1000; +/** Object-store listing and marker writes are metadata operations, not + * checkpoint transfers. Bound each tightly so the post-exec checkpoint + * reserve does not consume a full transfer timeout for a tiny request. */ +export const CHECKPOINT_METADATA_TIMEOUT_CAP_MS = 5_000; + +/** Worst-case time reserved after a successful session execute for the + * checkpoint pipeline: + * token throttle + guest pull + object listing + object upload + marker. + * The two large transfers receive the configured transfer timeout; the two + * metadata operations receive the smaller metadata cap. */ +export function checkpointPipelineBudgetMs( + launchTimeoutMs: number, + checkpointTimeoutMs: number, +): number { + const metadataTimeoutMs = Math.min( + checkpointTimeoutMs, + CHECKPOINT_METADATA_TIMEOUT_CAP_MS, + ); + return launchTimeoutMs + 2 * checkpointTimeoutMs + 2 * metadataTimeoutMs; +} + +/** BullMQ's `timestamp` is the enqueue time. Anchor the worker deadline to it + * so queueing consumes the same caller-visible JOB_TIMEOUT budget instead of + * granting a second full timeout after a delayed job finally starts. */ +export function jobDeadlineAtMs( + enqueuedAtMs: number | undefined, + timeoutMs: number, + nowMs: number = Date.now(), +): number { + return Number.isFinite(enqueuedAtMs) && (enqueuedAtMs as number) > 0 + ? (enqueuedAtMs as number) + timeoutMs + : nowMs + timeoutMs; +} + +export function parseArnList(raw: string | undefined): string[] | undefined { + if (raw == null) return undefined; + const entries = raw.split(',').map((entry) => entry.trim()).filter((entry) => entry.length > 0); + return entries.length > 0 ? entries : undefined; +} + +export interface LambdaMicrovmNumericConfig { + LAMBDA_MICROVM_PORT: number; + LAMBDA_MICROVM_MAX_DURATION_SECONDS: number; + LAMBDA_MICROVM_IDLE_SECONDS: number; + LAMBDA_MICROVM_SUSPEND_SECONDS: number; + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: number; + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: number; + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: number; + LAMBDA_MICROVM_LAUNCH_TPS: number; + LAMBDA_MICROVM_TOKEN_TPS: number; +} + +type LambdaMicrovmNumericEnv = Record; + +interface IntegerRange { + min: number; + max?: number; +} + +const lambdaMicrovmNumericRanges: Record = { + LAMBDA_MICROVM_PORT: { min: 1, max: 65_535 }, + LAMBDA_MICROVM_MAX_DURATION_SECONDS: { min: 1, max: 28_800 }, + LAMBDA_MICROVM_IDLE_SECONDS: { min: 60, max: 28_800 }, + LAMBDA_MICROVM_SUSPEND_SECONDS: { min: 0, max: 28_800 }, + /* Keep proxy credentials shorter than the AWS 60-minute maximum. */ + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: { min: 1, max: 900 }, + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: { min: 1 }, + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: { min: 1 }, + LAMBDA_MICROVM_LAUNCH_TPS: { min: 1 }, + LAMBDA_MICROVM_TOKEN_TPS: { min: 1 }, +}; + +const lambdaMicrovmNumericDefaults: LambdaMicrovmNumericConfig = { + LAMBDA_MICROVM_PORT: 8080, + LAMBDA_MICROVM_MAX_DURATION_SECONDS: 28_800, + /* 30min keeps a session's VM fully RUNNING (RAM + page cache live, ~0.3s + * follow-ups) across a realistic conversation gap before it suspends; + * 5min proved too aggressive — heavy libraries (chdb ~400MB) pay a + * 30-120s lazy rootfs re-read whenever the cache is lost. */ + LAMBDA_MICROVM_IDLE_SECONDS: 1_800, + LAMBDA_MICROVM_SUSPEND_SECONDS: 1_800, + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: 300, + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: 60_000, + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: 5_000, + LAMBDA_MICROVM_LAUNCH_TPS: 4, + LAMBDA_MICROVM_TOKEN_TPS: 8, +}; + +/** Parse configured values without `||` so an intentional zero survives long + * enough for the range validator to accept it where AWS does (suspend duration) + * and reject it everywhere else. */ +export function resolveLambdaMicrovmNumericConfig( + source: LambdaMicrovmNumericEnv, +): LambdaMicrovmNumericConfig { + const read = (name: keyof LambdaMicrovmNumericConfig): number => { + const raw = source[name]; + return raw == null || raw.trim() === '' ? lambdaMicrovmNumericDefaults[name] : Number(raw); + }; + return { + LAMBDA_MICROVM_PORT: read('LAMBDA_MICROVM_PORT'), + LAMBDA_MICROVM_MAX_DURATION_SECONDS: read('LAMBDA_MICROVM_MAX_DURATION_SECONDS'), + LAMBDA_MICROVM_IDLE_SECONDS: read('LAMBDA_MICROVM_IDLE_SECONDS'), + LAMBDA_MICROVM_SUSPEND_SECONDS: read('LAMBDA_MICROVM_SUSPEND_SECONDS'), + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: read('LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS'), + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: read('LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS'), + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: read('LAMBDA_MICROVM_HEALTH_TIMEOUT_MS'), + LAMBDA_MICROVM_LAUNCH_TPS: read('LAMBDA_MICROVM_LAUNCH_TPS'), + LAMBDA_MICROVM_TOKEN_TPS: read('LAMBDA_MICROVM_TOKEN_TPS'), + }; +} + +/** Returns the first invalid Lambda numeric setting for fail-fast startup. */ +export function lambdaMicrovmNumericConfigError( + config: LambdaMicrovmNumericConfig, +): string | undefined { + for (const name of Object.keys(lambdaMicrovmNumericRanges) as Array) { + const value = config[name]; + const { min, max } = lambdaMicrovmNumericRanges[name]; + if (!Number.isSafeInteger(value) || value < min || (max != null && value > max)) { + const range = max == null ? `at least ${min}` : `between ${min} and ${max}`; + return `${name} must be a whole number ${range}`; + } + } + return undefined; +} export function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, jobTimeoutMs: number): number { const defaultTtlSeconds = Math.max(1, Math.ceil((jobTimeoutMs + EGRESS_GRANT_GRACE_MS) / 1000)); @@ -63,6 +188,12 @@ export function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, return Math.max(1, Math.ceil(configuredTtlSeconds)); } +const lambdaMicrovmNumericConfig = resolveLambdaMicrovmNumericConfig(process.env); + +function configuredNumber(raw: string | undefined, fallback: number): number { + return raw == null || raw.trim() === '' ? fallback : Number(raw); +} + export const env = { PORT: process.env.SERVICE_PORT ?? 3112, LOCAL_MODE: process.env.LOCAL_MODE === 'true', @@ -142,6 +273,56 @@ export const env = { */ PTC_MODE: (process.env.PTC_MODE === 'blocking' ? 'blocking' : 'replay') as 'replay' | 'blocking', PTC_DEBUG: process.env.PTC_DEBUG === 'true', + /** + * Sandbox execution backend. + * - `http` (default): POST signed execute requests to SANDBOX_ENDPOINT + * (current Kubernetes/libkrun sandbox-runner). + * - `lambda-microvm`: AWS Lambda MicroVM backend. + */ + SANDBOX_BACKEND: (process.env.CODEAPI_SANDBOX_BACKEND === 'lambda-microvm' + ? 'lambda-microvm' + : 'http') as 'http' | 'lambda-microvm', + /** + * Runtime session affinity for stateful sandbox backends. + * - `stateless` (default): no runtime sessions; `runtime_session_hint` ignored. + * - `affinity`: stateful session reuse; contention surfaces as HTTP 409 + * rather than silently losing workspace state in a stateless execution. + * - `strict`: same serialized session semantics, and a session hint is + * required instead of degrading requests without one to stateless. + */ + RUNTIME_SESSION_MODE: (process.env.CODEAPI_RUNTIME_SESSION_MODE === 'affinity' + || process.env.CODEAPI_RUNTIME_SESSION_MODE === 'strict' + ? process.env.CODEAPI_RUNTIME_SESSION_MODE + : 'stateless') as 'stateless' | 'affinity' | 'strict', + RUNTIME_SESSION_LOCK_WAIT_MS: configuredNumber( + process.env.CODEAPI_RUNTIME_SESSION_LOCK_WAIT_MS, + 15_000, + ), + // Lambda MicroVM backend. Connector lists are comma-separated ARNs. + LAMBDA_MICROVM_IMAGE_ARN: process.env.LAMBDA_MICROVM_IMAGE_ARN ?? '', + LAMBDA_MICROVM_IMAGE_VERSION: process.env.LAMBDA_MICROVM_IMAGE_VERSION || undefined, + LAMBDA_MICROVM_EXECUTION_ROLE_ARN: process.env.LAMBDA_MICROVM_EXECUTION_ROLE_ARN || undefined, + /* Runtime VM stdout reaches CloudWatch only when RunMicrovm sends a logging + * config AND an executionRoleArn is set — pairs with the role above. */ + LAMBDA_MICROVM_LOG_GROUP: process.env.LAMBDA_MICROVM_LOG_GROUP || undefined, + LAMBDA_MICROVM_REGION: process.env.LAMBDA_MICROVM_REGION || undefined, + LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS: parseArnList(process.env.LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS), + LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS: parseArnList(process.env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS), + ...lambdaMicrovmNumericConfig, + /* CreateMicrovmAuthToken is minted per execute + per checkpoint; share a + * fleet-wide budget so concurrent warm-session executes queue instead of + * bursting past the AWS TPS limit. */ + LAMBDA_MICROVM_ALLOW_SHELL: process.env.LAMBDA_MICROVM_ALLOW_SHELL === 'true', + /* Session workspace checkpoints (effective only in affinity/strict modes). + * On by default so VM expiry/eviction recovery is automatic; the byte cap + * bounds tar size pulled from the VM and stored to S3. */ + SESSION_CHECKPOINTS: process.env.CODEAPI_SESSION_CHECKPOINTS !== 'false', + CHECKPOINT_MAX_BYTES: configuredNumber( + process.env.CODEAPI_CHECKPOINT_MAX_BYTES, + 512 * 1024 * 1024, + ), + CHECKPOINT_TIMEOUT_MS: configuredNumber(process.env.CODEAPI_CHECKPOINT_TIMEOUT_MS, 60_000), + CHECKPOINT_PREFIX: process.env.CODEAPI_CHECKPOINT_PREFIX ?? 'rtsx-checkpoints/', }; const default_run_memory_limit = 256 * 1024 * 1024; diff --git a/service/src/file-server.ts b/service/src/file-server.ts index dd89bd2..6d226c0 100644 --- a/service/src/file-server.ts +++ b/service/src/file-server.ts @@ -719,16 +719,22 @@ app.delete('/sessions/:session_id/objects/:fileId', async (req, res) => { } }); -const port = process.env.FILE_SERVER_PORT ?? 3000; +const port = Number(process.env.FILE_SERVER_PORT ?? 3000); +/** Optional bind address. Deployments where only the co-located service-api + * should reach the file server (push-model sandbox backends fetch input + * bytes server-side) set this to 127.0.0.1 so the object routes are never + * network-exposed. Unset preserves the historical all-interfaces bind. */ +const host = process.env.FILE_SERVER_HOST; let server: ReturnType | undefined; let shuttingDown = false; async function startServer(): Promise { try { await initializeStorage(); - server = app.listen(port, () => { - logger.info(`[${INSTANCE_ID}] Server running on port ${port}`); - }); + const onListen = () => { + logger.info(`[${INSTANCE_ID}] Server running on ${host ?? '*'}:${port}`); + }; + server = host ? app.listen(port, host, onListen) : app.listen(port, onListen); } catch (err) { logger.error('Critical: Could not initialize storage', { error: err }); process.exit(1); diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 6deec64..59a18a8 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -3,7 +3,7 @@ import type { Express } from 'express'; import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from './queue'; import { validateStartupAuthConfig } from './auth/startup'; import { env } from './config'; -import { validateApiHardenedConfig, validateWorkerHardenedConfig } from './secure-startup'; +import { validateApiHardenedConfig, validateSandboxBackendPolicy, validateWorkerHardenedConfig } from './secure-startup'; import logger from './logger'; import { shutdownTelemetry } from './telemetry'; @@ -75,6 +75,11 @@ function setupQueueListeners(queue: Queue, name: string): void { export async function startupApiOnly(): Promise { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); + /* No validateSandboxBackendPolicy() here: an API-only pod authenticates and + * enqueues jobs, it never constructs the Lambda backend or checkpoint store. + * Validating that policy would force worker-only config (LAMBDA_MICROVM_* and + * the MINIO_* checkpoint creds) into API pods just to boot. The worker and + * combined startups own that validation. */ await validateLifecycleAuthConfig(); // Set up queue listeners for monitoring (optional, for observability) @@ -92,6 +97,7 @@ export async function startupApiOnly(): Promise { export async function startupWorkerOnly(): Promise { logger.info('Starting Worker service...'); validateWorkerHardenedConfig(); + validateSandboxBackendPolicy(); // Dynamically import workers to start them const { pyWorker, otherWorker } = await import('./workers'); @@ -125,6 +131,7 @@ async function gracefulStartup(): Promise { logger.info('Starting up service (combined API + Workers)...'); validateApiHardenedConfig(); validateWorkerHardenedConfig(); + validateSandboxBackendPolicy(); await validateLifecycleAuthConfig(); try { diff --git a/service/src/metrics.ts b/service/src/metrics.ts index a6ce6cd..ba052dc 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -194,6 +194,57 @@ export const ptcReplayStaleCleanups = new Counter({ help: 'Stale executions reaped by the periodic cleanup sweep', }); +export const microvmLaunches = new Counter({ + name: 'codeapi_microvm_launches_total', + help: 'Lambda MicroVM launch attempts by outcome', + labelNames: ['outcome'] as const, +}); + +export const microvmLaunchDuration = new Histogram({ + name: 'codeapi_microvm_launch_duration_seconds', + help: 'Time from RunMicrovm to a healthy RUNNING MicroVM', + buckets: [0.5, 1, 2, 5, 10, 20, 40, 60], +}); + +export const microvmTerminations = new Counter({ + name: 'codeapi_microvm_terminations_total', + help: 'Lambda MicroVM terminations by reason', + labelNames: ['reason'] as const, +}); + +export const microvmThrottleEvents = new Counter({ + name: 'codeapi_microvm_throttle_events_total', + help: 'Control-plane throttle waits/errors by operation', + labelNames: ['op'] as const, +}); + +export const runtimeSessionLockContention = new Counter({ + name: 'codeapi_runtime_session_lock_contention_total', + help: 'Runtime session lock waits that timed out, by mode', + labelNames: ['mode'] as const, +}); + +export const microvmCheckpoints = new Counter({ + name: 'codeapi_microvm_checkpoints_total', + help: 'Session workspace checkpoint attempts by outcome', + labelNames: ['outcome'] as const, +}); + +export const microvmRestores = new Counter({ + name: 'codeapi_microvm_restores_total', + help: 'Session workspace restore attempts by outcome', + labelNames: ['outcome'] as const, +}); + +export const microvmCheckpointBytes = new Histogram({ + name: 'codeapi_microvm_checkpoint_bytes', + help: 'Size of stored session workspace checkpoints', + buckets: [ + 1024, 64 * 1024, 1024 * 1024, 16 * 1024 * 1024, 64 * 1024 * 1024, + 256 * 1024 * 1024, 512 * 1024 * 1024, + ], +}); + // -- Helpers for serving metrics -- export async function metricsHandler(_req: unknown, res: { set: (key: string, value: string) => void; send: (data: string) => void }): Promise { diff --git a/service/src/runtime-session/checkpoint-store.test.ts b/service/src/runtime-session/checkpoint-store.test.ts new file mode 100644 index 0000000..918fec9 --- /dev/null +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -0,0 +1,230 @@ +import { describe, expect, test } from 'bun:test'; +import * as fsp from 'fs/promises'; +import { + MemoryCheckpointStore, + MinioCheckpointStore, + CheckpointTooLargeError, + checkpointObjectKey, + checkpointPrefixFor, + resolveS3Endpoint, +} from './checkpoint-store'; + +const BIG = 1_000_000; + +async function readStored( + store: MemoryCheckpointStore, + runtimeSessionId: string, + objectKey?: string, +): Promise { + const artifact = await store.get(runtimeSessionId, BIG, objectKey); + if (!artifact) return null; + try { + return await fsp.readFile(artifact.path, 'utf8'); + } finally { + await artifact.cleanup(); + } +} + +describe('resolveS3Endpoint', () => { + const cases: Array<{ + name: string; + options: Parameters[0]; + expected: string; + }> = [ + { + name: 'keeps the HTTPS default for a scheme-qualified S3 endpoint', + options: { endpoint: 'https://s3.us-east-1.amazonaws.com' }, + expected: 'https://s3.us-east-1.amazonaws.com', + }, + { + name: 'keeps the HTTP default for a scheme-qualified endpoint', + options: { endpoint: 'http://storage.example.com' }, + expected: 'http://storage.example.com', + }, + { + name: 'uses port 9000 for a bare MinIO service name', + options: { endpoint: 'minio' }, + expected: 'http://minio:9000', + }, + { + name: 'uses SSL and port 9000 for a bare local endpoint', + options: { endpoint: 'localhost', useSsl: true }, + expected: 'https://localhost:9000', + }, + { + name: 'preserves a URL port when MINIO_PORT is absent', + options: { endpoint: 'http://minio.internal:9001' }, + expected: 'http://minio.internal:9001', + }, + { + name: 'lets MINIO_PORT override a port already present in the URL', + options: { + endpoint: 'https://s3.us-east-1.amazonaws.com:8443', + port: '9443', + }, + expected: 'https://s3.us-east-1.amazonaws.com:9443', + }, + { + name: 'lets MINIO_PORT override the bare-endpoint default', + options: { endpoint: 'minio', port: '9001' }, + expected: 'http://minio:9001', + }, + ]; + + for (const { name, options, expected } of cases) { + test(name, () => { + expect(resolveS3Endpoint(options)).toBe(expected); + }); + } +}); + +describe('checkpoint store', () => { + test('the S3-compatible store aborts a hung underlying operation at its hard deadline', async () => { + let observedAbort = false; + const client = { + send: (_command: unknown, options?: { abortSignal?: AbortSignal }) => ( + new Promise((_resolve, reject) => { + options?.abortSignal?.addEventListener('abort', () => { + observedAbort = true; + reject(options.abortSignal?.reason); + }, { once: true }); + }) + ), + }; + const store = new MinioCheckpointStore(client, { bucket: 'test', timeoutMs: 10 }); + await expect(store.latestSequence('rt_hung')).rejects.toThrow('checkpoint list timed out'); + expect(observedAbort).toBe(true); + }); + + test('object key is per session + sequence, zero-padded for lexical order', () => { + expect(checkpointObjectKey('rt_abc', 1)).toBe('rtsx-checkpoints/rt_abc/00000000000000000001.tar.gz'); + expect(checkpointPrefixFor('rt_abc')).toBe('rtsx-checkpoints/rt_abc/'); + /* zero-padding keeps lexical order == numeric sequence order across widths */ + expect(checkpointObjectKey('rt_abc', 2) > checkpointObjectKey('rt_abc', 1)).toBe(true); + expect(checkpointObjectKey('rt_abc', 10) > checkpointObjectKey('rt_abc', 9)).toBe(true); + expect(checkpointObjectKey('rt_xyz', 1)).not.toBe(checkpointObjectKey('rt_abc', 1)); + }); + + test('S3 pruning chunks deletes at the 1,000-object API limit', async () => { + const stale = Array.from( + { length: 1_001 }, + (_, index) => ({ Key: checkpointObjectKey('rt_many', index + 1) }), + ); + const deleteBatchSizes: number[] = []; + const client = { + send: (command: { constructor?: { name?: string }; input?: { + Delete?: { Objects?: unknown[] }; + } }) => { + if (command.constructor?.name === 'ListObjectsV2Command') { + return Promise.resolve({ Contents: stale, IsTruncated: false }); + } + if (command.constructor?.name === 'DeleteObjectsCommand') { + deleteBatchSizes.push(command.input?.Delete?.Objects?.length ?? 0); + return Promise.resolve({}); + } + throw new Error(`unexpected command ${command.constructor?.name}`); + }, + }; + const store = new MinioCheckpointStore(client, { bucket: 'test' }); + + await store.pruneOlderThan('rt_many', 2_000); + expect(deleteBatchSizes).toEqual([1_000, 1]); + }); + + test('S3 high-water selection ignores non-canonical keys in the prefix', async () => { + const client = { + send: () => Promise.resolve({ + Contents: [ + { Key: checkpointObjectKey('rt_keys', 5) }, + { Key: `${checkpointPrefixFor('rt_keys')}99999999999999999999-foreign.tar.gz` }, + ], + IsTruncated: false, + }), + }; + const store = new MinioCheckpointStore(client, { bucket: 'test' }); + expect(await store.latestSequence('rt_keys')).toBe(5); + }); + + test('memory store round-trips the highest-sequence bytes and copies defensively', async () => { + const store = new MemoryCheckpointStore(); + const original = Buffer.from('workspace-bytes'); + await store.put('rt_1', 1, original); + await store.commit('rt_1', 1); + + expect(await readStored(store, 'rt_1')).toBe('workspace-bytes'); + /* stored copy is independent of the caller's buffer */ + original.fill(0); + expect(await readStored(store, 'rt_1')).toBe('workspace-bytes'); + }); + + test('absent checkpoint returns null', async () => { + const store = new MemoryCheckpointStore(); + expect(await store.get('rt_missing', BIG)).toBeNull(); + }); + + test('get reads the highest sequence, not the last write', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_1', 2, Buffer.from('newer')); + await store.commit('rt_1', 2); + /* a crash-orphaned put from a LOWER sequence lands late but must NOT win */ + await store.put('rt_1', 1, Buffer.from('stale')); + expect(await readStored(store, 'rt_1')).toBe('newer'); + }); + + test('latestSequence reports the max retained sequence (for reset seeding)', async () => { + const store = new MemoryCheckpointStore(); + expect(await store.latestSequence('rt_1')).toBe(0); + await store.put('rt_1', 5, Buffer.from('v5')); + expect(await store.latestSequence('rt_1')).toBe(5); + /* other sessions don't leak into this one's max */ + await store.put('rt_2', 9, Buffer.from('v9')); + expect(await store.latestSequence('rt_1')).toBe(5); + }); + + test('uncommitted late uploads are ignored and pruning is explicit', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_1', 1, Buffer.from('v1')); + await store.commit('rt_1', 1); + await store.put('rt_1', 2, Buffer.from('v2')); + await store.commit('rt_1', 2); + await store.put('rt_1', 3, Buffer.from('uncommitted')); + await store.put('rt_other', 1, Buffer.from('other')); + expect(await readStored(store, 'rt_1')).toBe('v2'); + await store.pruneOlderThan('rt_1', 2); + const keys = [...store.objects.keys()]; + expect(keys).toContain(checkpointObjectKey('rt_1', 2)); + expect(keys).not.toContain(checkpointObjectKey('rt_1', 1)); + expect(keys).toContain(checkpointObjectKey('rt_1', 3)); + expect(keys).toContain(checkpointObjectKey('rt_other', 1)); + }); + + test('an exact fenced pointer wins over a newer recovery marker', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_1', 1, Buffer.from('pointed')); + await store.put('rt_1', 2, Buffer.from('newer-marker')); + await store.commit('rt_1', 2); + expect( + await readStored(store, 'rt_1', checkpointObjectKey('rt_1', 1)), + ).toBe('pointed'); + }); + + test('recovery skips an orphan newest marker and uses an older valid pair', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_1', 1, Buffer.from('recoverable')); + await store.commit('rt_1', 1); + store.committed.add(checkpointObjectKey('rt_1', 2)); + + expect(await readStored(store, 'rt_1')).toBe('recoverable'); + + await store.pruneOlderThan('rt_1', 3); + expect(store.committed.has(checkpointObjectKey('rt_1', 1))).toBe(false); + expect(store.committed.has(checkpointObjectKey('rt_1', 2))).toBe(false); + }); + + test('rejects a checkpoint larger than maxBytes', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_big', 1, Buffer.alloc(2048)); + await store.commit('rt_big', 1); + await expect(store.get('rt_big', 1024)).rejects.toBeInstanceOf(CheckpointTooLargeError); + }); +}); diff --git a/service/src/runtime-session/checkpoint-store.ts b/service/src/runtime-session/checkpoint-store.ts new file mode 100644 index 0000000..d971018 --- /dev/null +++ b/service/src/runtime-session/checkpoint-store.ts @@ -0,0 +1,473 @@ +import { + DeleteObjectsCommand, + GetObjectCommand, + HeadObjectCommand, + ListObjectsV2Command, + PutObjectCommand, + S3Client, +} from '@aws-sdk/client-s3'; +import * as fs from 'fs'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { Readable, Transform } from 'stream'; +import { pipeline } from 'stream/promises'; +import { env } from '../config'; + +interface S3SendClient { + send( + command: unknown, + options?: { abortSignal?: AbortSignal }, + ): Promise; +} + +interface CheckpointStoreClientOptions { + bucket?: string; + timeoutMs?: number; +} + +export function resolveS3Endpoint(options: { + endpoint: string; + port?: string; + useSsl?: boolean; +}): string { + const hasScheme = options.endpoint.includes('://'); + const protocol = options.useSsl === true ? 'https:' : 'http:'; + const parsed = new URL(hasScheme ? options.endpoint : `${protocol}//${options.endpoint}`); + const explicitPort = options.port?.trim(); + if (explicitPort) { + /* MINIO_PORT is an explicit deployment override, including when the URL + * already contains a different port. */ + parsed.port = explicitPort; + } else if (!parsed.port && !hasScheme) { + /* A bare endpoint is the local/MinIO shorthand and keeps MinIO's default. + * An absolute URL instead keeps its scheme default (80/443). */ + parsed.port = '9000'; + } + return parsed.toString().replace(/\/$/, ''); +} + +function s3Endpoint(): string { + return resolveS3Endpoint({ + endpoint: process.env.MINIO_ENDPOINT ?? 'localhost', + port: process.env.MINIO_PORT, + useSsl: process.env.MINIO_USE_SSL?.trim().toLowerCase() === 'true', + }); +} + +function createS3Client(): S3Client { + const accessKeyId = process.env.MINIO_ACCESS_KEY?.trim(); + const secretAccessKey = process.env.MINIO_SECRET_KEY?.trim(); + return new S3Client({ + endpoint: s3Endpoint(), + region: process.env.MINIO_REGION || process.env.AWS_REGION || 'us-east-1', + forcePathStyle: true, + ...(accessKeyId && secretAccessKey + ? { + credentials: { + accessKeyId, + secretAccessKey, + ...(process.env.MINIO_SESSION_TOKEN + ? { sessionToken: process.env.MINIO_SESSION_TOKEN } + : {}), + }, + } + : {}), + /* Checkpoint bytes already have an immutable, fenced object key. Avoid + * buffering streams for opportunistic SDK checksums; S3 TLS + SigV4 still + * protect transport integrity and the archive parser validates structure. */ + requestChecksumCalculation: 'WHEN_REQUIRED', + responseChecksumValidation: 'WHEN_REQUIRED', + }); +} + +/** + * Durable storage for session workspace checkpoints. Each checkpoint writes a + * distinct object under a per-session prefix keyed by a monotonic sequence + * (`/.tar.gz`). The session registry points + * at one exact immutable object; after that fenced pointer commits, a sibling + * `.committed` marker makes the object eligible for recovery if Redis is later + * lost. Restore never selects an uploaded-but-uncommitted object. The sequence + * comes from an atomic Redis reservation (see `allocateCheckpointSequence`) + * above both the counter and `latestSequence()`, so TTL reset recovery and + * concurrent stale holders cannot collide and ordering has no dependence on + * per-worker wall clocks. Older data and markers are best-effort pruned only + * after the new pointer and marker commit. + */ +export interface CheckpointStore { + /** Uploads immutable checkpoint data. This does not make it restorable or + * prune prior checkpoints; the caller commits it only after its fenced + * registry pointer succeeds. */ + put( + runtimeSessionId: string, + sequence: number, + data: CheckpointArtifact | Buffer, + ): Promise; + /** Writes the durable recovery marker for a checkpoint whose fenced pointer + * has already committed. */ + commit(runtimeSessionId: string, sequence: number): Promise; + /** Best-effort garbage collection after commit. */ + pruneOlderThan(runtimeSessionId: string, sequence: number): Promise; + /** Reads the exact `objectKey`, or the highest-sequence durably committed + * checkpoint when no pointer is available. `maxBytes` is enforced before and + * during the streamed download. Throws {@link CheckpointTooLargeError} when + * exceeded. */ + get( + runtimeSessionId: string, + maxBytes: number, + objectKey?: string, + ): Promise; + /** The highest sequence number retained under the session prefix (0 if none), + * used to re-seed a reset counter above already-stored objects. */ + latestSequence(runtimeSessionId: string): Promise; +} + +export class CheckpointTooLargeError extends Error {} + +export interface CheckpointArtifact { + path: string; + size: number; + cleanup(): Promise; +} + +export async function checkpointArtifactFromStream( + stream: Readable, + maxBytes: number, + prefix = 'codeapi-checkpoint-', +): Promise { + const dir = await fsp.mkdtemp(path.join(os.tmpdir(), prefix)); + const artifactPath = path.join(dir, 'workspace.tar.gz'); + let size = 0; + const limit = new Transform({ + transform(chunk: Buffer, _encoding, callback) { + size += chunk.length; + callback( + size > maxBytes + ? new CheckpointTooLargeError(`checkpoint exceeded maxBytes ${maxBytes}B`) + : null, + chunk, + ); + }, + }); + try { + await pipeline(stream, limit, fs.createWriteStream(artifactPath, { flags: 'wx', mode: 0o600 })); + return { + path: artifactPath, + size, + cleanup: () => fsp.rm(dir, { recursive: true, force: true }), + }; + } catch (error) { + await fsp.rm(dir, { recursive: true, force: true }).catch(() => {}); + throw error; + } +} + +async function checkpointArtifactFromBuffer(data: Buffer): Promise { + return checkpointArtifactFromStream(Readable.from([data]), data.length); +} + +/** Zero-padded so lexicographic key order matches numeric sequence order. */ +const SEQUENCE_WIDTH = 20; + +export function checkpointPrefixFor(runtimeSessionId: string): string { + return `${env.CHECKPOINT_PREFIX}${runtimeSessionId}/`; +} + +export function checkpointObjectKey(runtimeSessionId: string, sequence: number): string { + return `${checkpointPrefixFor(runtimeSessionId)}${String(sequence).padStart(SEQUENCE_WIDTH, '0')}.tar.gz`; +} + +function checkpointCommitKey(runtimeSessionId: string, sequence: number): string { + return `${checkpointObjectKey(runtimeSessionId, sequence)}.committed`; +} + +/** Parse the sequence back out of a full object key (inverse of the key fn). */ +function sequenceFromKey(key: string): number { + const base = key.slice(key.lastIndexOf('/') + 1).replace(/\.tar\.gz$/, ''); + const seq = Number.parseInt(base, 10); + return Number.isFinite(seq) ? seq : 0; +} + +/** S3/MinIO-backed store using the same MINIO_* envs as file-server. */ +export class MinioCheckpointStore implements CheckpointStore { + private readonly client: S3SendClient; + private readonly bucket: string; + private readonly timeoutMs: number; + + constructor(client?: S3SendClient, options: CheckpointStoreClientOptions = {}) { + this.client = client ?? createS3Client(); + /* `||` not `??`: an empty-string CODEAPI_CHECKPOINT_BUCKET must fall through + * to MINIO_BUCKET (startup validation accepts the config when MINIO_BUCKET + * is set, so `??` would otherwise select '' and fail every S3 op). */ + this.bucket = options.bucket + || process.env.CODEAPI_CHECKPOINT_BUCKET + || process.env.MINIO_BUCKET + || 'test-bucket'; + this.timeoutMs = options.timeoutMs ?? env.CHECKPOINT_TIMEOUT_MS; + } + + /** AWS SDK AbortSignal reaches the Node HTTP handler, so expiry destroys the + * underlying socket and upload/download stream instead of merely abandoning + * a still-running MinIO promise after the session lock has been released. */ + private async withDeadline( + label: string, + operation: (signal: AbortSignal) => Promise, + ): Promise { + const controller = new AbortController(); + const timer = setTimeout(() => { + controller.abort(new Error(`${label} timed out after ${this.timeoutMs}ms`)); + }, this.timeoutMs); + timer.unref?.(); + try { + return await operation(controller.signal); + } finally { + clearTimeout(timer); + } + } + + private send(label: string, command: unknown): Promise { + return this.withDeadline(label, async (abortSignal) => ( + await this.client.send(command, { abortSignal }) as T + )); + } + + async put( + runtimeSessionId: string, + sequence: number, + data: CheckpointArtifact | Buffer, + ): Promise { + const key = checkpointObjectKey(runtimeSessionId, sequence); + const size = Buffer.isBuffer(data) ? data.length : data.size; + await this.withDeadline('checkpoint upload', async (abortSignal) => { + const source = Buffer.isBuffer(data) + ? data + : fs.createReadStream(data.path, { signal: abortSignal }); + try { + await this.client.send(new PutObjectCommand({ + Bucket: this.bucket, + Key: key, + Body: source, + ContentLength: size, + ContentType: 'application/x-gtar', + }), { abortSignal }); + } finally { + if (!Buffer.isBuffer(source)) source.destroy(); + } + }); + } + + async commit(runtimeSessionId: string, sequence: number): Promise { + const marker = Buffer.from('committed\n'); + await this.send('checkpoint marker upload', new PutObjectCommand({ + Bucket: this.bucket, + Key: checkpointCommitKey(runtimeSessionId, sequence), + Body: marker, + ContentLength: marker.length, + ContentType: 'text/plain', + })); + } + + async pruneOlderThan(runtimeSessionId: string, sequence: number): Promise { + const keepKey = checkpointObjectKey(runtimeSessionId, sequence); + const stale = (await this.listKeys(runtimeSessionId)).filter(key => { + if (key.endsWith('.tar.gz')) return key < keepKey; + if (key.endsWith('.tar.gz.committed')) { + return key.slice(0, -'.committed'.length) < keepKey; + } + return false; + }); + for (let offset = 0; offset < stale.length; offset += 1_000) { + const chunk = stale.slice(offset, offset + 1_000); + const result = await this.send<{ Errors?: Array<{ Key?: string; Code?: string }> }>( + 'checkpoint prune', + new DeleteObjectsCommand({ + Bucket: this.bucket, + Delete: { Objects: chunk.map(Key => ({ Key })), Quiet: true }, + }), + ); + if (result.Errors?.length) { + throw new Error( + `Failed to prune checkpoint objects: ` + + result.Errors.map(item => `${item.Key ?? '?'}:${item.Code ?? '?'}`).join(', '), + ); + } + } + } + + async latestSequence(runtimeSessionId: string): Promise { + const key = await this.latestKey(runtimeSessionId); + return key ? sequenceFromKey(key) : 0; + } + + async get( + runtimeSessionId: string, + maxBytes: number, + objectKey?: string, + ): Promise { + const key = objectKey ?? await this.latestCommittedKey(runtimeSessionId); + if (!key) return null; + const prefix = checkpointPrefixFor(runtimeSessionId); + if (!key.startsWith(prefix) || !key.endsWith('.tar.gz')) { + throw new Error('Checkpoint pointer is outside the runtime session prefix'); + } + let size: number; + try { + const stat = await this.send<{ ContentLength?: number }>( + 'checkpoint stat', + new HeadObjectCommand({ Bucket: this.bucket, Key: key }), + ); + size = stat.ContentLength ?? 0; + } catch (error) { + const code = (error as { code?: string; name?: string })?.code + ?? (error as { name?: string })?.name; + const status = (error as { $metadata?: { httpStatusCode?: number } })?.$metadata?.httpStatusCode; + if (code === 'NoSuchKey' || code === 'NotFound' || status === 404) { + if (objectKey) throw new Error(`Committed checkpoint object is missing: ${key}`); + return null; + } + throw error; + } + /* Reject before downloading, then enforce again while streaming in case the + * object changes or a compatible store reports an inaccurate size. */ + if (size > maxBytes) { + throw new CheckpointTooLargeError(`checkpoint ${size}B exceeds maxBytes ${maxBytes}B`); + } + return this.withDeadline('checkpoint download', async (abortSignal) => { + const result = await this.client.send( + new GetObjectCommand({ Bucket: this.bucket, Key: key }), + { abortSignal }, + ) as { Body?: AsyncIterable }; + if (!result.Body) throw new Error('Checkpoint download response omitted Body'); + return checkpointArtifactFromStream(Readable.from(result.Body), maxBytes); + }); + } + + private async listKeys(runtimeSessionId: string): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + return this.withDeadline('checkpoint list', async (abortSignal) => { + const keys: string[] = []; + let ContinuationToken: string | undefined; + do { + const page = await this.client.send(new ListObjectsV2Command({ + Bucket: this.bucket, + Prefix: prefix, + ContinuationToken, + }), { abortSignal }) as { + Contents?: Array<{ Key?: string }>; + IsTruncated?: boolean; + NextContinuationToken?: string; + }; + for (const item of page.Contents ?? []) { + if (item.Key) keys.push(item.Key); + } + ContinuationToken = page.IsTruncated ? page.NextContinuationToken : undefined; + if (page.IsTruncated && !ContinuationToken) { + throw new Error('Checkpoint list response omitted continuation token'); + } + } while (ContinuationToken); + return keys; + }); + } + + private async latestKey(runtimeSessionId: string): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + const keys = (await this.listKeys(runtimeSessionId)).filter( + key => key.startsWith(prefix) + && /^\d{20}\.tar\.gz$/.test(key.slice(prefix.length)), + ); + if (keys.length === 0) return null; + return keys.reduce((max, key) => (key > max ? key : max)); + } + + private async latestCommittedKey(runtimeSessionId: string): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + const keys = await this.listKeys(runtimeSessionId); + const available = new Set(keys); + const markers = keys + .filter(key => key.startsWith(prefix) + && /^\d{20}\.tar\.gz\.committed$/.test(key.slice(prefix.length))); + if (markers.length === 0) return null; + markers.sort().reverse(); + for (const marker of markers) { + const dataKey = marker.slice(0, -'.committed'.length); + if (available.has(dataKey)) return dataKey; + } + return null; + } +} + +/** In-memory store for bun tests. Keyed by full object key so latest-selection + * and pruning mirror the S3-backed store. */ +export class MemoryCheckpointStore implements CheckpointStore { + readonly objects = new Map(); + readonly committed = new Set(); + + async put( + runtimeSessionId: string, + sequence: number, + data: CheckpointArtifact | Buffer, + ): Promise { + const key = checkpointObjectKey(runtimeSessionId, sequence); + this.objects.set(key, Buffer.isBuffer(data) ? Buffer.from(data) : await fsp.readFile(data.path)); + } + + async commit(runtimeSessionId: string, sequence: number): Promise { + this.committed.add(checkpointObjectKey(runtimeSessionId, sequence)); + } + + async pruneOlderThan(runtimeSessionId: string, sequence: number): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + const keep = checkpointObjectKey(runtimeSessionId, sequence); + for (const existing of this.objects.keys()) { + if (existing.startsWith(prefix) && existing < keep) { + this.objects.delete(existing); + this.committed.delete(existing); + } + } + for (const existing of this.committed) { + if (existing.startsWith(prefix) && existing < keep) { + this.committed.delete(existing); + } + } + } + + async latestSequence(runtimeSessionId: string): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + let max = 0; + for (const key of this.objects.keys()) { + if (key.startsWith(prefix)) max = Math.max(max, sequenceFromKey(key)); + } + return max; + } + + async get( + runtimeSessionId: string, + maxBytes: number, + objectKey?: string, + ): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + let latest = objectKey; + if (!latest) { + for (const key of this.committed) { + if ( + key.startsWith(prefix) + && this.objects.has(key) + && (latest === undefined || key > latest) + ) latest = key; + } + } + if (latest === undefined) return null; + if (!latest.startsWith(prefix) || !latest.endsWith('.tar.gz')) { + throw new Error('Checkpoint pointer is outside the runtime session prefix'); + } + const data = this.objects.get(latest); + if (!data) { + if (objectKey) throw new Error(`Committed checkpoint object is missing: ${latest}`); + return null; + } + if (data.length > maxBytes) { + throw new CheckpointTooLargeError(`checkpoint ${data.length}B exceeds maxBytes ${maxBytes}B`); + } + return checkpointArtifactFromBuffer(Buffer.from(data)); + } +} diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts new file mode 100644 index 0000000..4cb85dd --- /dev/null +++ b/service/src/runtime-session/checkpoint.ts @@ -0,0 +1,386 @@ +import axios from 'axios'; +import * as fs from 'fs'; +import { Readable } from 'stream'; +import type { MicrovmAuthToken } from './lambda-client'; +import { microvmPortHeaders } from './lambda-client'; +import { + CheckpointTooLargeError, + checkpointArtifactFromStream, + type CheckpointArtifact, + type CheckpointStore, +} from './checkpoint-store'; +import { + acquireRuntimeSessionLock, + allocateCheckpointSequence, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + writeRuntimeSessionRecord, +} from './registry'; +import { checkpointObjectKey } from './checkpoint-store'; +import { microvmCheckpoints, microvmRestores, microvmCheckpointBytes } from '../metrics'; +import { CHECKPOINT_METADATA_TIMEOUT_CAP_MS } from '../config'; +import logger from '../logger'; + +/** Reject if `promise` doesn't settle within `ms`, so a stalled metadata leg + * cannot hold the session lock. The production S3-compatible store separately + * attaches an AbortSignal deadline to every request and streamed transfer, + * which destroys the underlying transport when its hard bound expires. */ +function withTimeout( + promise: Promise, + ms: number, + label: string, + onLateValue?: (value: T) => void | Promise, +): Promise { + return new Promise((resolve, reject) => { + let timedOut = false; + const timer = setTimeout(() => { + timedOut = true; + reject(new Error(`${label} timed out after ${ms}ms`)); + }, ms); + promise.then( + (value) => { + clearTimeout(timer); + if (timedOut) { + if (onLateValue) void Promise.resolve(onLateValue(value)).catch(() => {}); + return; + } + resolve(value); + }, + (err) => { clearTimeout(timer); reject(err); }, + ); + }); +} + +/** + * Auto-checkpoint orchestration. The workspace only mutates during an + * execute, and executes serialize on the session lock — so a lock-guarded + * checkpoint after each exec yields complete, tear-free coverage: if a newer + * exec already holds the lock we skip, and that exec's own post-checkpoint + * covers our changes. Restore runs in-path on relaunch, before the first + * execute on the fresh VM. A missed post-exec checkpoint is non-fatal and + * leaves the prior durable pointer intact. Restore failures fail closed and + * recycle the fresh VM rather than silently running against an empty or + * partially restored workspace. + */ + +export interface CheckpointConfig { + port: number; + authTokenTtlSeconds: number; + maxBytes: number; + timeoutMs: number; +} + +/* Opts the runner into session mode for checkpoint/restore. These run before + * the first /execute on a relaunched VM, so the runner has nothing bound yet in + * the hookless design; without this header the handler 409s. Case-insensitive, + * matches the runner's `x-runtime-session-id`. */ +const RUNTIME_SESSION_ID_HEADER = 'X-Runtime-Session-Id'; + +export async function pullCheckpoint( + args: { + mintToken: () => Promise; + endpointBase: string; + runtimeSessionId: string; + signal?: AbortSignal; + }, + config: CheckpointConfig, +): Promise { + const token = await args.mintToken(); + const response = await axios.get(`${args.endpointBase}/api/v2/session/checkpoint`, { + headers: { + [token.headerName]: token.token, + ...microvmPortHeaders(config.port), + [RUNTIME_SESSION_ID_HEADER]: args.runtimeSessionId, + }, + responseType: 'stream', + timeout: config.timeoutMs, + signal: args.signal, + }); + const announced = Number(response.headers['content-length']); + if (Number.isFinite(announced) && announced > config.maxBytes) { + response.data.destroy(); + throw new CheckpointTooLargeError( + `checkpoint ${announced}B exceeds maxBytes ${config.maxBytes}B`, + ); + } + return checkpointArtifactFromStream(response.data, config.maxBytes); +} + +export async function pushRestore( + args: { + mintToken: () => Promise; + endpointBase: string; + runtimeSessionId: string; + signal?: AbortSignal; + }, + data: CheckpointArtifact, + config: CheckpointConfig, +): Promise { + const token = await args.mintToken(); + await axios.post(`${args.endpointBase}/api/v2/session/restore`, fs.createReadStream(data.path), { + headers: { + [token.headerName]: token.token, + ...microvmPortHeaders(config.port), + [RUNTIME_SESSION_ID_HEADER]: args.runtimeSessionId, + 'Content-Type': 'application/x-gtar', + 'Content-Length': String(data.size), + }, + maxBodyLength: config.maxBytes, + timeout: config.timeoutMs, + signal: args.signal, + }); +} + +/** + * Asks the VM which of `refs` its input cache is missing. Dedupe lives here, + * not in Redis: control-plane state can be lost with a recycle, while the VM + * cannot be wrong about what it holds. + */ +export async function probeInputs( + args: { mintToken: () => Promise; endpointBase: string; signal?: AbortSignal }, + refs: Array<{ cache_key: string }>, + config: CheckpointConfig, +): Promise> { + const token = await args.mintToken(); + const response = await axios.post<{ missing?: Array<{ cache_key: string }> }>( + `${args.endpointBase}/api/v2/session/inputs/probe`, + { refs }, + { + headers: { + [token.headerName]: token.token, + ...microvmPortHeaders(config.port), + 'Content-Type': 'application/json', + }, + timeout: config.timeoutMs, + signal: args.signal, + }, + ); + const missing = response.data?.missing; + if (!Array.isArray(missing)) { + throw new Error('Runner returned an invalid input probe response'); + } + const requested = new Set(refs.map(ref => ref.cache_key)); + const seen = new Set(); + for (const ref of missing) { + if ( + typeof ref?.cache_key !== 'string' || + !/^[0-9a-f]{64}$/.test(ref.cache_key) || + !requested.has(ref.cache_key) || + seen.has(ref.cache_key) + ) { + throw new Error('Runner returned an invalid input probe response'); + } + seen.add(ref.cache_key); + } + return missing; +} + +/** Pushes a digest-named input batch into the VM's runner-local cache. Never + * touches the sandbox workspace — priming does that, from the cache. */ +export async function pushInputs( + args: { mintToken: () => Promise; endpointBase: string; signal?: AbortSignal }, + data: NodeJS.ReadableStream | Buffer | (() => NodeJS.ReadableStream), + config: CheckpointConfig, + contentLength?: number, + expandedBytes?: number, +): Promise { + const token = await args.mintToken(); + /* Build file streams only after token minting succeeds. An eagerly-created + * fs.ReadStream opens on a later tick; if minting rejects, caller cleanup can + * unlink the archive first and the stream then emits an unhandled ENOENT. */ + const requestData = typeof data === 'function' ? data() : data; + await axios.post(`${args.endpointBase}/api/v2/session/inputs`, requestData, { + headers: { + [token.headerName]: token.token, + ...microvmPortHeaders(config.port), + 'Content-Type': 'application/x-gtar', + ...(contentLength === undefined ? {} : { 'Content-Length': String(contentLength) }), + ...(expandedBytes === undefined + ? {} + : { 'X-CodeAPI-Input-Expanded-Bytes': String(expandedBytes) }), + }, + maxBodyLength: Math.max(config.maxBytes, contentLength ?? 0), + timeout: config.timeoutMs, + signal: args.signal, + }); +} + +/** + * Checkpoint the session workspace: pull the tar from the still-warm VM, + * store it, and record the pointer under the lock (fenced write). Pass + * `lockToken` to reuse a lock already held (the post-exec path); omit it for + * a standalone checkpoint (e.g. a pre-deadline sweep), which takes a single + * non-blocking lock — a busy lock means a newer exec is running and its own + * post-checkpoint will cover this one. + */ +export async function checkpointSession(args: { + mintToken: (microvmId: string) => Promise; + store: CheckpointStore; + runtimeSessionId: string; + config: CheckpointConfig; + normalizeEndpoint: (endpoint: string) => string; + lockToken?: string; + signal?: AbortSignal; +}): Promise<'stored' | 'skipped_busy' | 'skipped_state' | 'failed'> { + const heldToken = args.lockToken; + const lockToken = heldToken ?? await acquireRuntimeSessionLock(args.runtimeSessionId); + let data: CheckpointArtifact | undefined; + if (!lockToken) { + microvmCheckpoints.inc({ outcome: 'skipped_busy' }); + return 'skipped_busy'; + } + try { + const record = await readRuntimeSessionRecord(args.runtimeSessionId); + if (!record || record.state !== 'RUNNING' || !record.microvm_id || !record.endpoint) { + microvmCheckpoints.inc({ outcome: 'skipped_state' }); + return 'skipped_state'; + } + const microvmId = record.microvm_id; + data = await pullCheckpoint({ + mintToken: () => args.mintToken(microvmId), + endpointBase: args.normalizeEndpoint(record.endpoint), + runtimeSessionId: args.runtimeSessionId, + signal: args.signal, + }, args.config); + /* Read the durable high-water mark before reserving a sequence. The Redis + * reservation atomically advances above max(counter, retainedMax), so even + * a stale holder that resumes after lock expiry receives a distinct object + * key and cannot overwrite the newer holder's committed checkpoint. */ + const retainedMax = await withTimeout( + args.store.latestSequence(args.runtimeSessionId), + Math.min(args.config.timeoutMs, CHECKPOINT_METADATA_TIMEOUT_CAP_MS), + 'checkpoint store.latestSequence', + ); + const sequence = await allocateCheckpointSequence(args.runtimeSessionId, retainedMax); + /* Upload immutable data first. If this times out or we lose the lock before + * the pointer CAS, it is only an uncommitted orphan: restore ignores it. */ + await withTimeout( + args.store.put(args.runtimeSessionId, sequence, data), + args.config.timeoutMs, + 'checkpoint store.put', + ); + + /* Commit the exact pointer under the session fence only after the object + * exists. This prevents a crash/timeout from publishing a missing object. */ + const persisted = await writeRuntimeSessionRecord({ + ...record, + workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId, sequence), + checkpointed_at: Date.now(), + }, lockToken); + if (!persisted) { + microvmCheckpoints.inc({ outcome: 'skipped_busy' }); + return 'skipped_busy'; + } + + /* The Redis pointer is authoritative while present. A durable commit + * marker lets a later restore recover after Redis TTL/loss without ever + * selecting an uncommitted late upload. Marker/GC failures must not undo + * the already-committed pointer. */ + let markerCommitted = false; + await withTimeout( + args.store.commit(args.runtimeSessionId, sequence), + Math.min(args.config.timeoutMs, CHECKPOINT_METADATA_TIMEOUT_CAP_MS), + 'checkpoint store.commit', + ).then(() => { + markerCommitted = true; + }).catch(error => { + logger.warn('Checkpoint commit marker failed; Redis pointer remains authoritative', { + runtimeSessionId: args.runtimeSessionId, + error: error instanceof Error ? error.message : String(error), + }); + }); + /* Never prune the previous durable recovery point unless the new marker + * exists. The Redis pointer can outlive a marker failure, but Redis loss + * must still fall back to the prior committed pair. */ + if (markerCommitted) { + void args.store.pruneOlderThan(args.runtimeSessionId, sequence).catch(error => { + logger.warn('Checkpoint garbage collection failed', { + runtimeSessionId: args.runtimeSessionId, + error: error instanceof Error ? error.message : String(error), + }); + }); + } + microvmCheckpointBytes.observe(data.size); + microvmCheckpoints.inc({ outcome: 'stored' }); + return 'stored'; + } catch (error) { + microvmCheckpoints.inc({ outcome: 'failed' }); + logger.warn('Session checkpoint failed', { + runtimeSessionId: args.runtimeSessionId, + error: error instanceof Error ? error.message : String(error), + }); + return 'failed'; + } finally { + await data?.cleanup().catch(() => {}); + /* Only release a lock we acquired here. */ + if (!heldToken) await releaseRuntimeSessionLock(args.runtimeSessionId, lockToken); + } +} + +/** Relaunch restore: caller holds the session lock and the VM is RUNNING. */ +export async function restoreSession(args: { + mintToken: (microvmId: string) => Promise; + store: CheckpointStore; + runtimeSessionId: string; + microvmId: string; + endpointBase: string; + config: CheckpointConfig; + signal?: AbortSignal; + checkpointKey?: string; +}): Promise<'restored' | 'absent' | 'fetch_failed' | 'push_failed'> { + /* The store enforces `maxBytes` before and during its streamed download, so an + * oversized/stray checkpoint cannot consume unbounded worker memory or disk. + * Bound the fetch too — a stalled S3/MinIO get would otherwise hold the + * session lock through the whole relaunch and time the request out. */ + let data: CheckpointArtifact | null; + try { + data = await withTimeout( + args.store.get(args.runtimeSessionId, args.config.maxBytes, args.checkpointKey), + args.config.timeoutMs, + 'checkpoint store.get', + late => late?.cleanup(), + ); + } catch (error) { + /* Fetch failed before the runner was touched — the workspace is still the + * clean fresh-VM one, but a prior checkpoint pointer means running there + * would silently lose session state. The caller fails closed and recycles + * this VM. */ + microvmRestores.inc({ outcome: 'failed' }); + logger.warn('Checkpoint fetch failed; refusing to run with an empty workspace', { + runtimeSessionId: args.runtimeSessionId, + error: error instanceof Error ? error.message : String(error), + }); + return 'fetch_failed'; + } + if (!data) { + microvmRestores.inc({ outcome: 'absent' }); + return 'absent'; + } + try { + await pushRestore({ + mintToken: () => args.mintToken(args.microvmId), + endpointBase: args.endpointBase, + runtimeSessionId: args.runtimeSessionId, + signal: args.signal, + }, data, args.config); + microvmRestores.inc({ outcome: 'restored' }); + logger.info('Session workspace restored from checkpoint', { + runtimeSessionId: args.runtimeSessionId, + bytes: data.size, + }); + return 'restored'; + } catch (error) { + /* Push failed AFTER the runner may have started extracting/chowning/wiping + * the workspace. That cleanup runs async past our client abort, so the + * workspace is possibly-partial — the caller must recycle the VM rather than + * execute against it. */ + microvmRestores.inc({ outcome: 'failed' }); + logger.warn('Checkpoint push-restore failed; the VM workspace may be partial', { + runtimeSessionId: args.runtimeSessionId, + error: error instanceof Error ? error.message : String(error), + }); + return 'push_failed'; + } finally { + await data.cleanup().catch(() => {}); + } +} diff --git a/service/src/runtime-session/files.test.ts b/service/src/runtime-session/files.test.ts new file mode 100644 index 0000000..b88a25d --- /dev/null +++ b/service/src/runtime-session/files.test.ts @@ -0,0 +1,167 @@ +import { afterAll, beforeAll, describe, expect, test } from 'bun:test'; +import { spawnSync } from 'child_process'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { + SESSION_INPUTS_MAX_COUNT, + buildInputBatch, + inputCacheKey, + sessionFileRefs, +} from './files'; + +let server: ReturnType; + +beforeAll(() => { + server = Bun.serve({ + port: 0, + fetch(req) { + const pathname = new URL(req.url).pathname; + if (pathname.endsWith('/objects/fmissing')) { + return new Response('not found', { status: 404 }); + } + if (pathname.endsWith('/objects/fsource')) { + return new Response('upstream failure', { status: 503 }); + } + const readOnly = pathname.includes('/objects/ro'); + const headers: Record = { 'X-Original-Filename': 'server-name.txt' }; + if (readOnly) headers['X-Read-Only'] = 'true'; + return new Response('0123456789', { status: 200, headers }); + }, + }); +}); + +afterAll(() => { + server.stop(true); +}); + +const opts = (overrides: Partial<{ maxBytes: number; signal: AbortSignal }> = {}) => ({ + timeoutMs: 5_000, + maxBytes: 1024 * 1024, + fileServerUrl: `http://localhost:${server.port}`, + ...overrides, +}); + +const ref = (n: number | string) => ({ + id: `f${n}`, + storage_session_id: 's1', + name: `file-${n}.txt`, + cache_key: inputCacheKey('s1', `f${n}`), +}); + +/** Lists member names of a produced batch. */ +async function membersOf(archive: string): Promise { + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'batch-check-')); + try { + spawnSync('tar', ['-xzf', archive, '-C', tmp]); + return (await fsp.readdir(tmp)).sort(); + } finally { + await fsp.rm(tmp, { recursive: true, force: true }); + } +} + +/** + * Cross-component contract — see the twin assertion in the runner's + * session-inputs.test.ts. Members are named here and looked up there, so a + * silent divergence makes every delivery land under names the runner can never + * find. That happened (NUL vs space separator) and only live execution caught + * it; this vector is what makes it a test failure instead. + */ +const GOLDEN_KEY_SID_1_FILE_1 = 'a995f1e7977466c5636419d21582e0b44420c44d2d7e2660b13aa4d4b4667d90'; + +describe('inputCacheKey', () => { + test('matches the digest the runner resolves cache entries with', () => { + expect(inputCacheKey('sid-1', 'file-1')).toBe(GOLDEN_KEY_SID_1_FILE_1); + }); +}); + +describe('sessionFileRefs', () => { + test('collapses the same object requested under multiple names', () => { + /* Identity is (storage session, id): one delivery, and priming writes it + * to each requested path. Two entries would push the same bytes twice. */ + const refs = sessionFileRefs([ + { id: 'f1', storage_session_id: 's1', name: 'a.csv' }, + { id: 'f1', storage_session_id: 's1', name: 'copy/a.csv' }, + { id: 'f2', storage_session_id: 's1', name: 'b.csv' }, + { name: 'inline.py', content: 'print(1)' }, + ]); + expect(refs.map((r) => r.id)).toEqual(['f1', 'f2']); + }); +}); + +describe('buildInputBatch', () => { + test('packs digest-named members with metadata the runner can serve', async () => { + const batch = await buildInputBatch([ref(1)], opts()); + try { + expect(batch?.count).toBe(1); + expect(batch?.expandedSize).toBe( + 10 + Buffer.byteLength(JSON.stringify({ readOnly: false })), + ); + const key = inputCacheKey('s1', 'f1'); + expect(await membersOf(batch!.path)).toEqual([key, `${key}.json`].sort()); + } finally { + await batch?.cleanup(); + } + }); + + test('carries only object-level metadata (read-only), never a name', async () => { + const batch = await buildInputBatch( + [{ + id: 'ro', + storage_session_id: 's1', + name: 'skill.md', + cache_key: inputCacheKey('s1', 'ro'), + }], + opts(), + ); + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'meta-check-')); + try { + spawnSync('tar', ['-xzf', batch!.path, '-C', tmp]); + const meta = JSON.parse( + await fsp.readFile(path.join(tmp, `${inputCacheKey('s1', 'ro')}.json`), 'utf8'), + ); + /* A name here would override every requesting ref's destination. */ + expect(meta).toEqual({ readOnly: true }); + } finally { + await batch?.cleanup(); + await fsp.rm(tmp, { recursive: true, force: true }); + } + }); + + test('rejects deliveries above the object-count cap before any fetch', async () => { + const refs = Array.from({ length: SESSION_INPUTS_MAX_COUNT + 1 }, (_, i) => ref(i)); + await expect(buildInputBatch(refs, opts())).rejects.toMatchObject({ + code: 'SESSION_INPUT_TOO_LARGE', + }); + }); + + test('enforces a CUMULATIVE byte budget, not just per-object size', async () => { + await expect( + buildInputBatch([ref(1), ref(2), ref(3)], opts({ maxBytes: 25 })), + ).rejects.toMatchObject({ + code: 'SESSION_INPUT_TOO_LARGE', + }); + }); + + test('honors an aborted signal instead of consuming disk and bandwidth', async () => { + const controller = new AbortController(); + controller.abort(); + await expect( + buildInputBatch([ref(1)], opts({ signal: controller.signal })), + ).rejects.toMatchObject({ + code: 'SESSION_INPUT_ABORTED', + }); + }); + + test('classifies a rejected object reference as unavailable input', async () => { + await expect(buildInputBatch([ref('missing')], opts())).rejects.toMatchObject({ + code: 'SESSION_INPUT_UNAVAILABLE', + }); + }); + + test('classifies file-server failures separately from caller input errors', async () => { + await expect(buildInputBatch([ref('source')], opts())).rejects.toMatchObject({ + code: 'SESSION_INPUT_SOURCE_FAILED', + }); + }); +}); diff --git a/service/src/runtime-session/files.ts b/service/src/runtime-session/files.ts new file mode 100644 index 0000000..0a0a6ae --- /dev/null +++ b/service/src/runtime-session/files.ts @@ -0,0 +1,275 @@ +import axios from 'axios'; +import { spawn } from 'child_process'; +import * as crypto from 'crypto'; +import * as fs from 'fs'; +import * as fsp from 'fs/promises'; +import * as os from 'os'; +import * as path from 'path'; +import { Readable, Transform } from 'stream'; +import { pipeline } from 'stream/promises'; +import type * as t from '../types'; +import { internalServiceHeaders } from '../internal-service-auth'; +import { getAxiosErrorDetails } from '../utils'; +import { env } from '../config'; +import logger from '../logger'; + +/** + * Input delivery for sandbox backends whose guest cannot reach the file + * server (a MicroVM's only egress is the public internet, so the runner's + * pull-based priming has nothing reachable to pull from). + * + * The control plane fetches the authorized objects locally and pushes the + * bytes into a runner-local cache keyed by (storage session, object id). The + * runner's EXISTING priming path then resolves refs from that cache instead of + * over HTTP — so the sandbox workspace keeps exactly one writer, and pushed + * inputs inherit its identity, ownership, read-only, symlink-safety, priming + * and modification-detection semantics unchanged. + * + * Two consequences worth stating, because earlier designs got them wrong: + * - Nothing here decides what the workspace should contain. Re-pushing an + * object can never revert a sandbox edit, because the push does not touch + * the workspace at all. + * - Dedupe is asked of the VM (`probe`), not tracked in Redis. Control-plane + * state can be lost with a recycle; the VM's own cache cannot lie about + * what it holds. + */ + +export const SESSION_INPUTS_MAX_COUNT = 256; + +/** + * Stable input-delivery failures. These codes survive the worker/BullMQ + * boundary and let the public router distinguish a bad/oversized input set + * from a broken MicroVM without exposing file-server or archive details. + */ +export type SessionFilesErrorCode = + | 'SESSION_INPUT_TOO_LARGE' + | 'SESSION_INPUT_UNAVAILABLE' + | 'SESSION_INPUT_SOURCE_FAILED' + | 'SESSION_INPUT_PREPARATION_FAILED' + | 'SESSION_INPUT_ABORTED'; + +export class SessionFilesError extends Error { + constructor( + public readonly code: SessionFilesErrorCode, + message: string, + ) { + super(message); + this.name = 'SessionFilesError'; + } +} + +export interface SessionFileRef { + id: string; + storage_session_id: string; + name: string; + cache_key: string; +} + +/** Mirrors the runner's `inputCacheKey` (api/src/session-inputs.ts): both ends + * ship in the same image, so the digest is a hard-coded contract. */ +export function inputCacheKey(storageSessionId: string, id: string): string { + return crypto.createHash('sha256').update(`${storageSessionId}\u0000${id}`, 'utf8').digest('hex'); +} + +/** The by-reference subset of the payload's files (inline `content` entries + * need no delivery — the runner writes those itself). */ +export function sessionFileRefs(files: t.PayloadBody['files'] | undefined): SessionFileRef[] { + if (!files?.length) return []; + const refs: SessionFileRef[] = []; + const seen = new Set(); + for (const file of files) { + if (!('id' in file) || !file.id || !file.storage_session_id || !file.name) continue; + /* Identity is (storage session, id) — the same object requested under two + * names is ONE delivery; the runner writes it to each requested path from + * the payload during priming. */ + const key = inputCacheKey(file.storage_session_id, file.id); + if (seen.has(key)) continue; + seen.add(key); + refs.push({ + id: file.id, + storage_session_id: file.storage_session_id, + name: file.name, + cache_key: file.input_cache_key ?? key, + }); + } + return refs; +} + +export interface InputBatch { + path: string; + /** Exact uncompressed bytes in data files plus metadata sidecars. */ + expandedSize: number; + size: number; + count: number; + cleanup(): Promise; +} + +/** + * Fetches the given objects from the file server and packs them into the + * digest-named batch the runner's cache endpoint accepts. Throws + * {@link SessionFilesError} on a failed fetch or a blown budget — a silently + * missing input is the failure mode this module exists to prevent. + */ +export async function buildInputBatch( + refs: SessionFileRef[], + opts: { timeoutMs: number; maxBytes: number; fileServerUrl?: string; signal?: AbortSignal }, +): Promise { + if (refs.length === 0) return undefined; + if (refs.length > SESSION_INPUTS_MAX_COUNT) { + throw new SessionFilesError( + 'SESSION_INPUT_TOO_LARGE', + `Session delivery of ${refs.length} objects exceeds the ${SESSION_INPUTS_MAX_COUNT} limit`, + ); + } + const baseUrl = (opts.fileServerUrl ?? env.FILE_SERVER_URL).replace(/\/$/, ''); + const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'codeapi-inputs-')); + const objects = path.join(tmp, 'objects'); + const archive = path.join(tmp, 'inputs.tar.gz'); + try { + await fsp.mkdir(objects, { mode: 0o700 }); + let totalBytes = 0; + for (const ref of refs) { + if (opts.signal?.aborted) { + throw new SessionFilesError( + 'SESSION_INPUT_ABORTED', + 'Session input delivery aborted', + ); + } + const key = ref.cache_key; + const fetched = await fetchFileObjectToPath( + baseUrl, + ref, + path.join(objects, key), + { ...opts, remainingBytes: opts.maxBytes - totalBytes }, + ); + totalBytes += fetched.bytes; + const metadata = Buffer.from(JSON.stringify({ readOnly: fetched.readOnly })); + totalBytes += metadata.length; + if (totalBytes > opts.maxBytes) { + throw new SessionFilesError( + 'SESSION_INPUT_TOO_LARGE', + `Session inputs exceed the ${opts.maxBytes}-byte budget`, + ); + } + await fsp.writeFile( + path.join(objects, `${key}.json`), + /* Only object-level facts travel with the object; the destination + * name belongs to each requesting ref (see CachedInputMeta). */ + metadata, + ); + } + await tarDirectory(objects, archive, opts.signal); + const stat = await fsp.lstat(archive); + await fsp.rm(objects, { recursive: true, force: true }); + return { + path: archive, + expandedSize: totalBytes, + size: stat.size, + count: refs.length, + cleanup: () => fsp.rm(tmp, { recursive: true, force: true }), + }; + } catch (error) { + await fsp.rm(tmp, { recursive: true, force: true }).catch(() => {}); + if (error instanceof SessionFilesError) throw error; + logger.error('Failed to prepare session input batch:', getAxiosErrorDetails(error)); + throw new SessionFilesError( + 'SESSION_INPUT_PREPARATION_FAILED', + 'Failed to prepare session input batch', + ); + } +} + +async function fetchFileObjectToPath( + baseUrl: string, + ref: SessionFileRef, + destination: string, + opts: { + timeoutMs: number; + maxBytes: number; + remainingBytes: number; + signal?: AbortSignal; + }, +): Promise<{ bytes: number; readOnly: boolean }> { + const url = `${baseUrl}/sessions/${encodeURIComponent(ref.storage_session_id)}/objects/${encodeURIComponent(ref.id)}`; + try { + const response = await axios.get(url, { + headers: internalServiceHeaders(), + responseType: 'stream', + timeout: opts.timeoutMs, + signal: opts.signal, + }); + const announced = Number(response.headers['content-length']); + if (Number.isFinite(announced) && announced > opts.remainingBytes) { + response.data.destroy(); + throw new SessionFilesError( + 'SESSION_INPUT_TOO_LARGE', + `Session inputs exceed the ${opts.maxBytes}-byte budget`, + ); + } + let bytes = 0; + const limit = new Transform({ + transform(chunk: Buffer, _encoding, callback) { + bytes += chunk.length; + callback( + bytes > opts.remainingBytes + ? new SessionFilesError( + 'SESSION_INPUT_TOO_LARGE', + `Session inputs exceed the ${opts.maxBytes}-byte budget`, + ) + : null, + chunk, + ); + }, + }); + await pipeline(response.data, limit, fs.createWriteStream(destination, { flags: 'wx', mode: 0o600 })); + const readOnly = String(response.headers['x-read-only'] ?? '').toLowerCase() === 'true'; + return { bytes, readOnly }; + } catch (error) { + await fsp.rm(destination, { force: true }).catch(() => {}); + if (error instanceof SessionFilesError) throw error; + if (opts.signal?.aborted) { + throw new SessionFilesError( + 'SESSION_INPUT_ABORTED', + 'Session input delivery aborted', + ); + } + /* Sanitized details only: a raw axios error carries the request config — + * including the internal service token header — straight into the logs. */ + logger.error(`Failed to fetch session input ${ref.id}:`, getAxiosErrorDetails(error)); + const status = axios.isAxiosError(error) ? error.response?.status : undefined; + throw new SessionFilesError( + status != null && status >= 400 && status < 500 + ? 'SESSION_INPUT_UNAVAILABLE' + : 'SESSION_INPUT_SOURCE_FAILED', + `Failed to fetch input ${ref.name} from file server`, + ); + } +} + +function tarDirectory(root: string, archive: string, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const tar = spawn('tar', ['-czf', archive, '-C', root, '.'], { + stdio: ['ignore', 'ignore', 'pipe'], + env: { ...process.env, COPYFILE_DISABLE: '1' }, + signal, + }); + const errChunks: Buffer[] = []; + tar.stderr.on('data', (chunk: Buffer) => errChunks.push(chunk)); + tar.on('error', (error) => { + reject(new SessionFilesError( + 'SESSION_INPUT_PREPARATION_FAILED', + `Failed to start inputs tar: ${error.message}`, + )); + }); + tar.on('close', (code) => { + if (code !== 0) { + reject(new SessionFilesError( + 'SESSION_INPUT_PREPARATION_FAILED', + `inputs tar exited ${code}: ${Buffer.concat(errChunks).toString()}`, + )); + return; + } + resolve(); + }); + }); +} diff --git a/service/src/runtime-session/id.test.ts b/service/src/runtime-session/id.test.ts new file mode 100644 index 0000000..4e8aff4 --- /dev/null +++ b/service/src/runtime-session/id.test.ts @@ -0,0 +1,127 @@ +import { describe, expect, test } from 'bun:test'; +import { + RUNTIME_SESSION_HINT_MAX_LENGTH, + RuntimeSessionHintError, + deriveRuntimeSessionId, + resolveRuntimeSessionIdForExecRequest, + resolveRuntimeSessionIdForRequest, + validateRuntimeSessionHint, +} from './id'; + +const BASE = { storageNamespace: 'tenant-a', canonicalUserId: 'user-1' }; + +describe('deriveRuntimeSessionId', () => { + test('is deterministic and shaped rt_<40 hex>', () => { + const first = deriveRuntimeSessionId({ ...BASE, hint: 'conv-1' }); + const second = deriveRuntimeSessionId({ ...BASE, hint: 'conv-1' }); + expect(first).toBe(second); + expect(first).toMatch(/^rt_[0-9a-f]{40}$/); + }); + + test('separates tenants, users, and hints', () => { + const base = deriveRuntimeSessionId({ ...BASE, hint: 'conv-1' }); + expect(deriveRuntimeSessionId({ storageNamespace: 'tenant-b', canonicalUserId: 'user-1', hint: 'conv-1' })).not.toBe(base); + expect(deriveRuntimeSessionId({ storageNamespace: 'tenant-a', canonicalUserId: 'user-2', hint: 'conv-1' })).not.toBe(base); + expect(deriveRuntimeSessionId({ ...BASE, hint: 'conv-2' })).not.toBe(base); + }); + + test('absent hint maps to a stable per-user default session', () => { + expect(deriveRuntimeSessionId(BASE)).toBe(deriveRuntimeSessionId({ ...BASE, hint: undefined })); + expect(deriveRuntimeSessionId(BASE)).not.toBe(deriveRuntimeSessionId({ ...BASE, hint: 'conv-1' })); + }); + + test('field boundaries cannot be forged across namespace/user/hint', () => { + const a = deriveRuntimeSessionId({ storageNamespace: 'ten', canonicalUserId: 'ant-user', hint: 'h' }); + const b = deriveRuntimeSessionId({ storageNamespace: 'ten-ant', canonicalUserId: 'user', hint: 'h' }); + expect(a).not.toBe(b); + }); +}); + +describe('validateRuntimeSessionHint', () => { + test('passes through valid hints and treats absent/empty as undefined', () => { + expect(validateRuntimeSessionHint('conv_123.a:b-c')).toBe('conv_123.a:b-c'); + expect(validateRuntimeSessionHint(undefined)).toBeUndefined(); + expect(validateRuntimeSessionHint(null)).toBeUndefined(); + expect(validateRuntimeSessionHint('')).toBeUndefined(); + }); + + test('rejects non-strings, oversize, and forbidden characters', () => { + expect(() => validateRuntimeSessionHint(42)).toThrow(RuntimeSessionHintError); + expect(() => validateRuntimeSessionHint({})).toThrow(RuntimeSessionHintError); + expect(() => validateRuntimeSessionHint('x'.repeat(RUNTIME_SESSION_HINT_MAX_LENGTH + 1))).toThrow('at most'); + expect(() => validateRuntimeSessionHint('has space')).toThrow('may only contain'); + expect(() => validateRuntimeSessionHint('emoji🙂')).toThrow('may only contain'); + expect(validateRuntimeSessionHint('x'.repeat(RUNTIME_SESSION_HINT_MAX_LENGTH))).toHaveLength(RUNTIME_SESSION_HINT_MAX_LENGTH); + }); +}); + +describe('resolveRuntimeSessionIdForRequest', () => { + test('stateless mode never derives an id, even with a hint', () => { + expect(resolveRuntimeSessionIdForRequest({ mode: 'stateless', ...BASE, hint: 'conv-1' })).toBeUndefined(); + }); + + test('affinity and strict modes derive the same id for the same inputs', () => { + const affinity = resolveRuntimeSessionIdForRequest({ mode: 'affinity', ...BASE, hint: 'conv-1' }); + const strict = resolveRuntimeSessionIdForRequest({ mode: 'strict', ...BASE, hint: 'conv-1' }); + expect(affinity).toBeDefined(); + expect(affinity).toBe(strict as string); + }); + + /* A hintless request must never land on a session: deriving from the + * default hint would silently share one persistent per-user workspace + * across every hintless conversation. */ + test('affinity mode without a hint degrades to stateless', () => { + expect(resolveRuntimeSessionIdForRequest({ mode: 'affinity', ...BASE })).toBeUndefined(); + expect(resolveRuntimeSessionIdForRequest({ mode: 'affinity', ...BASE, hint: '' })).toBeUndefined(); + }); + + test('strict mode without a hint is rejected', () => { + expect(() => resolveRuntimeSessionIdForRequest({ mode: 'strict', ...BASE })).toThrow( + RuntimeSessionHintError, + ); + }); +}); + +describe('resolveRuntimeSessionIdForExecRequest', () => { + test('stateless requests continue to ignore malformed hints', () => { + expect(resolveRuntimeSessionIdForExecRequest({ + mode: 'stateless', + ...BASE, + runtimeSessionHint: { ignored: true }, + isSynthetic: false, + })).toBeUndefined(); + }); + + test('strict-mode synthetic requests remain stateless without a hint', () => { + expect(resolveRuntimeSessionIdForExecRequest({ + mode: 'strict', + ...BASE, + runtimeSessionHint: undefined, + isSynthetic: true, + })).toBeUndefined(); + }); + + test('synthetic requests bypass validation for an ignored malformed hint', () => { + expect(resolveRuntimeSessionIdForExecRequest({ + mode: 'strict', + ...BASE, + runtimeSessionHint: { forged: true }, + isSynthetic: true, + })).toBeUndefined(); + }); + + test('strict-mode ordinary requests still validate and require a hint', () => { + expect(() => resolveRuntimeSessionIdForExecRequest({ + mode: 'strict', + ...BASE, + runtimeSessionHint: undefined, + isSynthetic: false, + })).toThrow('runtime_session_hint is required in strict mode'); + expect(() => resolveRuntimeSessionIdForExecRequest({ + mode: 'strict', + ...BASE, + runtimeSessionHint: { forged: true }, + isSynthetic: false, + })).toThrow('runtime_session_hint must be a string'); + }); +}); diff --git a/service/src/runtime-session/id.ts b/service/src/runtime-session/id.ts new file mode 100644 index 0000000..2a316c3 --- /dev/null +++ b/service/src/runtime-session/id.ts @@ -0,0 +1,95 @@ +import { createHash } from 'crypto'; + +export const RUNTIME_SESSION_HINT_MAX_LENGTH = 128; +const RUNTIME_SESSION_HINT_PATTERN = /^[A-Za-z0-9._:-]+$/; +const DEFAULT_HINT = 'default'; + +export class RuntimeSessionHintError extends Error { + readonly status = 400; + constructor(message: string) { + super(message); + this.name = 'RuntimeSessionHintError'; + } +} + +/** Normalizes the client-supplied hint: absent/empty ⇒ undefined, malformed ⇒ 400. */ +export function validateRuntimeSessionHint(hint: unknown): string | undefined { + if (hint == null) return undefined; + if (typeof hint !== 'string') { + throw new RuntimeSessionHintError('runtime_session_hint must be a string'); + } + if (hint.length === 0) return undefined; + if (hint.length > RUNTIME_SESSION_HINT_MAX_LENGTH) { + throw new RuntimeSessionHintError( + `runtime_session_hint must be at most ${RUNTIME_SESSION_HINT_MAX_LENGTH} characters`, + ); + } + if (!RUNTIME_SESSION_HINT_PATTERN.test(hint)) { + throw new RuntimeSessionHintError( + 'runtime_session_hint may only contain letters, digits, ".", "_", ":", and "-"', + ); + } + return hint; +} + +/** + * Server-derived runtime session identity. The namespace and user come from + * `getExecutionIdentity(req)` — never the client — so a hint can never + * collide across tenants or users. The hint only partitions sessions within + * one (tenant, user) scope. + */ +export function deriveRuntimeSessionId(args: { + storageNamespace: string; + canonicalUserId: string; + hint?: string; +}): string { + const material = `${args.storageNamespace}\u0000${args.canonicalUserId}\u0000${args.hint ?? DEFAULT_HINT}`; + return `rt_${createHash('sha256').update(material, 'utf8').digest('hex').slice(0, 40)}`; +} + +/** + * Router-side gate: stateless mode never derives a runtime session, and a + * request WITHOUT a hint never lands on a session either — deriving from the + * default hint would silently share one persistent per-user workspace across + * every hintless conversation, contradicting the caller's toggle-off + * expectation. Affinity degrades the hintless request to a stateless one-shot; + * strict mode rejects it, since the caller asked for guaranteed session + * semantics it failed to identify. + */ +export function resolveRuntimeSessionIdForRequest(args: { + mode: 'stateless' | 'affinity' | 'strict'; + storageNamespace: string; + canonicalUserId: string; + hint?: string; +}): string | undefined { + if (args.mode === 'stateless') return undefined; + if (args.hint == null || args.hint === '') { + if (args.mode === 'strict') { + throw new RuntimeSessionHintError('runtime_session_hint is required in strict mode'); + } + return undefined; + } + return deriveRuntimeSessionId(args); +} + +/** + * `/exec`-specific resolution boundary. Synthetic probes are deliberately + * sessionless, so they must bypass both strict-mode hint requirements and + * validation of a hint that will never be consumed. Stateless mode has the + * same ignore-don't-validate contract. + */ +export function resolveRuntimeSessionIdForExecRequest(args: { + mode: 'stateless' | 'affinity' | 'strict'; + storageNamespace: string; + canonicalUserId: string; + runtimeSessionHint: unknown; + isSynthetic: boolean; +}): string | undefined { + if (args.isSynthetic || args.mode === 'stateless') return undefined; + return resolveRuntimeSessionIdForRequest({ + mode: args.mode, + storageNamespace: args.storageNamespace, + canonicalUserId: args.canonicalUserId, + hint: validateRuntimeSessionHint(args.runtimeSessionHint), + }); +} diff --git a/service/src/runtime-session/input-delivery.test.ts b/service/src/runtime-session/input-delivery.test.ts new file mode 100644 index 0000000..52e3f19 --- /dev/null +++ b/service/src/runtime-session/input-delivery.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, test } from 'bun:test'; +import type { ExecutionManifestClaims } from '../execution-manifest'; +import { prepareSandboxEgress } from '../egress-grant'; +import type * as t from '../types'; +import { prepareInputDelivery } from './input-delivery'; +import { inputCacheKey } from './files'; + +const SECRET = 'test-egress-secret-32-bytes-minimum'; + +function claims(): ExecutionManifestClaims { + return { + v: 1, + exec_id: 'exec_123', + tenant_id: 'tenant_abc', + user_id: 'user_123', + session_key: 'tenant:tenant_abc:user:user_123', + input_files: [{ id: 'raw-file', session_id: 'raw-session', name: 'data.csv' }], + read_sessions: ['raw-session'], + output_session_id: 'raw-output', + max_upload_bytes: 1024, + max_output_files: 10, + max_requests: 100, + iat: 100, + exp: 300, + principal_source: 'librechat', + auth_context_hash: 'hash_123', + }; +} + +function payload(): t.PayloadBody { + return { + language: 'python', + version: '3.14.4', + files: [ + { id: 'raw-file', storage_session_id: 'raw-session', name: 'data.csv' }, + { id: 'raw-file', storage_session_id: 'raw-session', name: 'copy/data.csv' }, + { name: 'main.py', content: 'print(1)' }, + ], + }; +} + +describe('prepareInputDelivery', () => { + test('keeps raw fetch refs control-plane-only and gives fresh grants one stable cache key', () => { + const authorized = payload(); + const firstMasked = prepareSandboxEgress({ + payload: authorized, + claims: claims(), + grantId: 'grant_1', + secret: SECRET, + }).payload; + const secondMasked = prepareSandboxEgress({ + payload: authorized, + claims: claims(), + grantId: 'grant_2', + secret: SECRET, + }).payload; + + const first = prepareInputDelivery(authorized, firstMasked); + const second = prepareInputDelivery(authorized, secondMasked); + const stable = inputCacheKey('raw-session', 'raw-file'); + const firstRef = first.payload.files[0] as t.PayloadFileRef; + const secondRef = second.payload.files[0] as t.PayloadFileRef; + + expect(firstRef.id).not.toBe(secondRef.id); + expect(firstRef.input_cache_key).toBe(stable); + expect(secondRef.input_cache_key).toBe(stable); + expect(first.refs).toEqual([{ + id: 'raw-file', + storage_session_id: 'raw-session', + name: 'data.csv', + cache_key: stable, + }]); + expect(JSON.stringify(first.payload)).not.toContain('raw-file'); + expect(JSON.stringify(first.payload)).not.toContain('raw-session'); + }); + + test('rejects a gateway response that changes file ordering or shape', () => { + const authorized = payload(); + expect(() => prepareInputDelivery(authorized, { + ...authorized, + files: [...authorized.files].reverse(), + })).toThrow('shape or destination'); + }); +}); diff --git a/service/src/runtime-session/input-delivery.ts b/service/src/runtime-session/input-delivery.ts new file mode 100644 index 0000000..6d4432a --- /dev/null +++ b/service/src/runtime-session/input-delivery.ts @@ -0,0 +1,51 @@ +import type * as t from '../types'; +import type { SandboxInputDeliveryRef } from '../sandbox-backend/types'; +import { inputCacheKey } from './files'; + +function isFileRef(file: t.PayloadBody['files'][number]): file is t.PayloadFileRef { + return 'id' in file && typeof file.id === 'string'; +} + +/** + * Joins the authorized payload to its sandbox-visible counterpart by position. + * Hardened egress preserves file ordering while replacing ref identifiers with + * per-grant handles. The returned raw refs never cross the sandbox boundary; + * only the stable digest is added to the manifest-bound execute body. + */ +export function prepareInputDelivery( + authorized: t.PayloadBody, + sandboxVisible: t.PayloadBody, +): { payload: t.PayloadBody; refs: SandboxInputDeliveryRef[] } { + if (authorized.files.length !== sandboxVisible.files.length) { + throw new Error('Sandbox egress changed the input file count'); + } + + const refs: SandboxInputDeliveryRef[] = []; + const seen = new Set(); + const files = sandboxVisible.files.map((sandboxFile, index) => { + const authorizedFile = authorized.files[index]; + const authorizedIsRef = isFileRef(authorizedFile); + const sandboxIsRef = isFileRef(sandboxFile); + if (authorizedIsRef !== sandboxIsRef || authorizedFile.name !== sandboxFile.name) { + throw new Error(`Sandbox egress changed files[${index}] shape or destination`); + } + if (!authorizedIsRef || !sandboxIsRef) return { ...sandboxFile }; + + const cacheKey = inputCacheKey( + authorizedFile.storage_session_id, + authorizedFile.id, + ); + if (!seen.has(cacheKey)) { + seen.add(cacheKey); + refs.push({ + id: authorizedFile.id, + storage_session_id: authorizedFile.storage_session_id, + name: authorizedFile.name, + cache_key: cacheKey, + }); + } + return { ...sandboxFile, input_cache_key: cacheKey }; + }); + + return { payload: { ...sandboxVisible, files }, refs }; +} diff --git a/service/src/runtime-session/job-policy.test.ts b/service/src/runtime-session/job-policy.test.ts new file mode 100644 index 0000000..21771df --- /dev/null +++ b/service/src/runtime-session/job-policy.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, test } from 'bun:test'; +import { + PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, + resolveRuntimeSessionIdForJob, +} from './job-policy'; + +describe('resolveRuntimeSessionIdForJob', () => { + test('keeps explicitly exempt programmatic jobs stateless in strict mode', () => { + expect(resolveRuntimeSessionIdForJob({ + mode: 'strict', + runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, + isSynthetic: false, + })).toBeUndefined(); + }); + + test('the programmatic exemption wins over an accidentally supplied session id', () => { + expect(resolveRuntimeSessionIdForJob({ + mode: 'affinity', + runtimeSessionId: 'rt_should_not_be_used', + runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, + isSynthetic: false, + })).toBeUndefined(); + }); + + test('strict ordinary jobs still require a runtime session id', () => { + expect(() => resolveRuntimeSessionIdForJob({ + mode: 'strict', + isSynthetic: false, + })).toThrow('strict runtime session mode requires a runtimeSessionId on the job'); + }); + + test('synthetic jobs retain their existing strict-mode exemption', () => { + expect(resolveRuntimeSessionIdForJob({ + mode: 'strict', + isSynthetic: true, + })).toBeUndefined(); + }); +}); diff --git a/service/src/runtime-session/job-policy.ts b/service/src/runtime-session/job-policy.ts new file mode 100644 index 0000000..cc0635f --- /dev/null +++ b/service/src/runtime-session/job-policy.ts @@ -0,0 +1,27 @@ +import type { RuntimeSessionExemption } from '../types'; + +export const PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION: RuntimeSessionExemption = 'programmatic'; + +/** + * Resolve the session identity a worker is allowed to pass to its backend. + * Programmatic/replay jobs intentionally remain stateless until PTC can be + * bound to a warm session safely; the explicit queue marker makes that + * exception distinguishable from an ordinary `/exec` producer bug. + */ +export function resolveRuntimeSessionIdForJob(args: { + mode: 'stateless' | 'affinity' | 'strict'; + runtimeSessionId?: string; + runtimeSessionExemption?: RuntimeSessionExemption; + isSynthetic: boolean; +}): string | undefined { + if ( + args.isSynthetic + || args.runtimeSessionExemption === PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION + ) { + return undefined; + } + if (args.mode === 'strict' && !args.runtimeSessionId) { + throw new Error('strict runtime session mode requires a runtimeSessionId on the job'); + } + return args.runtimeSessionId; +} diff --git a/service/src/runtime-session/lambda-client-aws.test.ts b/service/src/runtime-session/lambda-client-aws.test.ts new file mode 100644 index 0000000..e8ca6b9 --- /dev/null +++ b/service/src/runtime-session/lambda-client-aws.test.ts @@ -0,0 +1,270 @@ +import { describe, expect, test } from 'bun:test'; +import { AwsLambdaMicrovmClient, type MicrovmCommandSender } from './lambda-client-aws'; +import { LambdaMicrovmApiError, MICROVM_AUTH_HEADER } from './lambda-client'; + +type SentCommand = { constructor: { name: string }; input: Record }; + +function stubSender(responses: unknown[]): { sender: MicrovmCommandSender; sent: SentCommand[] } { + const sent: SentCommand[] = []; + return { + sent, + sender: { + send(command: unknown): Promise { + sent.push(command as SentCommand); + const next = responses.shift(); + if (next instanceof Error) return Promise.reject(next); + return Promise.resolve(next); + }, + }, + }; +} + +function namedError(name: string): Error { + const error = new Error(`${name} raised`); + error.name = name; + return error; +} + +describe('AwsLambdaMicrovmClient command mapping', () => { + test('runMicrovm maps args onto RunMicrovmCommand input and normalizes the response', async () => { + const startedAt = new Date('2026-07-05T00:00:00Z'); + const { sender, sent } = stubSender([{ + microvmId: 'mvm-1', + state: 'PENDING', + endpoint: 'https://mvm-1.on.aws', + imageArn: 'arn:aws:lambda:us-east-2:1:microvm-image/codeapi', + imageVersion: '7', + maximumDurationInSeconds: 28_800, + startedAt, + }]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + + const description = await client.runMicrovm({ + imageIdentifier: 'arn:aws:lambda:us-east-2:1:microvm-image/codeapi', + imageVersion: '7', + executionRoleArn: 'arn:aws:iam::1:role/codeapi-microvm', + ingressConnectorArns: ['arn:ingress'], + egressConnectorArns: ['arn:egress'], + maximumDurationSeconds: 28_800, + idlePolicy: { maxIdleSeconds: 300, suspendedSeconds: 1_800, autoResume: true }, + runHookPayload: '{"runtime_session_id":"rt_x"}', + clientToken: 'launch-rt_x-7', + }); + + expect(sent[0].constructor.name).toBe('RunMicrovmCommand'); + expect(sent[0].input).toEqual({ + imageIdentifier: 'arn:aws:lambda:us-east-2:1:microvm-image/codeapi', + imageVersion: '7', + executionRoleArn: 'arn:aws:iam::1:role/codeapi-microvm', + ingressNetworkConnectors: ['arn:ingress'], + egressNetworkConnectors: ['arn:egress'], + maximumDurationInSeconds: 28_800, + idlePolicy: { + maxIdleDurationSeconds: 300, + suspendedDurationSeconds: 1_800, + autoResumeEnabled: true, + }, + runHookPayload: '{"runtime_session_id":"rt_x"}', + clientToken: 'launch-rt_x-7', + }); + expect(description).toEqual({ + microvmId: 'mvm-1', + state: 'PENDING', + endpoint: 'https://mvm-1.on.aws', + imageArn: 'arn:aws:lambda:us-east-2:1:microvm-image/codeapi', + imageVersion: '7', + maximumDurationSeconds: 28_800, + startedAtMs: startedAt.getTime(), + stateReason: undefined, + }); + }); + + test('lifecycle commands address the VM via microvmIdentifier', async () => { + const { sender, sent } = stubSender([ + { microvmId: 'mvm-1', state: 'RUNNING' }, + {}, + {}, + { microvmId: 'mvm-1', state: 'RUNNING' }, + {}, + ]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + + await client.getMicrovm('mvm-1'); + await client.suspendMicrovm('mvm-1'); + await client.resumeMicrovm('mvm-1'); + await client.terminateMicrovm('mvm-1'); + + expect(sent.map((command) => command.constructor.name)).toEqual([ + 'GetMicrovmCommand', + 'SuspendMicrovmCommand', + 'ResumeMicrovmCommand', + 'GetMicrovmCommand', + 'TerminateMicrovmCommand', + ]); + for (const command of sent) { + expect(command.input).toEqual({ microvmIdentifier: 'mvm-1' }); + } + }); + + test('createMicrovmAuthToken clamps TTL to whole minutes and reads the header map', async () => { + const { sender, sent } = stubSender([ + { authToken: { [MICROVM_AUTH_HEADER]: 'proxy-token-1' } }, + ]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + + const token = await client.createMicrovmAuthToken({ microvmId: 'mvm-1', port: 8080, ttlSeconds: 300 }); + + expect(sent[0].constructor.name).toBe('CreateMicrovmAuthTokenCommand'); + expect(sent[0].input).toEqual({ + microvmIdentifier: 'mvm-1', + expirationInMinutes: 5, + allowedPorts: [{ port: 8080 }], + }); + expect(token.headerName).toBe(MICROVM_AUTH_HEADER); + expect(token.token).toBe('proxy-token-1'); + expect(token.expiresAtMs).toBeGreaterThan(Date.now()); + }); + + test('token TTL clamps to the 1..60 minute API bounds', async () => { + const { sender, sent } = stubSender([ + { authToken: { [MICROVM_AUTH_HEADER]: 't1' } }, + { authToken: { [MICROVM_AUTH_HEADER]: 't2' } }, + ]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + + await client.createMicrovmAuthToken({ microvmId: 'mvm-1', port: 8080, ttlSeconds: 10 }); + await client.createMicrovmAuthToken({ microvmId: 'mvm-1', port: 8080, ttlSeconds: 86_400 }); + + expect((sent[0].input as { expirationInMinutes: number }).expirationInMinutes).toBe(1); + expect((sent[1].input as { expirationInMinutes: number }).expirationInMinutes).toBe(60); + }); + + test('missing token entry in the response surfaces as an API error', async () => { + const { sender } = stubSender([{ authToken: {} }]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + expect(client.createMicrovmAuthToken({ microvmId: 'mvm-1', port: 8080, ttlSeconds: 300 })) + .rejects.toThrow(`missing ${MICROVM_AUTH_HEADER}`); + }); + + test('a hanging SDK request is aborted by the client request deadline', async () => { + const sender: MicrovmCommandSender = { + send(_command, options): Promise { + return new Promise((_resolve, reject) => { + options?.abortSignal?.addEventListener( + 'abort', + () => reject(options.abortSignal?.reason ?? new Error('aborted')), + { once: true }, + ); + }); + }, + }; + const client = new AwsLambdaMicrovmClient({ client: sender, requestTimeoutMs: 10 }); + + await expect(client.getMicrovm('mvm-hung')).rejects.toMatchObject({ + operation: 'GetMicrovm', + kind: 'other', + }); + }); + + test('passes a caller abort signal through to the SDK request', async () => { + let observedSignal: AbortSignal | undefined; + const sender: MicrovmCommandSender = { + send(_command, options): Promise { + observedSignal = options?.abortSignal; + return new Promise((_resolve, reject) => { + observedSignal?.addEventListener( + 'abort', + () => reject(observedSignal?.reason ?? new Error('aborted')), + { once: true }, + ); + }); + }, + }; + const client = new AwsLambdaMicrovmClient({ client: sender, requestTimeoutMs: 5_000 }); + const controller = new AbortController(); + const pending = client.createMicrovmAuthToken({ + microvmId: 'mvm-hung', + port: 8080, + ttlSeconds: 300, + }, controller.signal); + controller.abort(new Error('job timed out')); + + await expect(pending).rejects.toMatchObject({ operation: 'CreateMicrovmAuthToken' }); + expect(observedSignal?.aborted).toBe(true); + }); + + test('replays the same client token after an ambiguous abort to recover the VM id', async () => { + const sent: Array<{ command: SentCommand; signal?: AbortSignal }> = []; + const sender: MicrovmCommandSender = { + send(command, options): Promise { + sent.push({ command: command as SentCommand, signal: options?.abortSignal }); + if (sent.length === 1) { + return new Promise((_resolve, reject) => { + options?.abortSignal?.addEventListener( + 'abort', + () => reject(options.abortSignal?.reason ?? new Error('aborted')), + { once: true }, + ); + }); + } + return Promise.resolve({ + microvmId: 'mvm-recovered', + state: 'RUNNING', + endpoint: 'https://mvm-recovered.example', + }); + }, + }; + const client = new AwsLambdaMicrovmClient({ client: sender, requestTimeoutMs: 5_000 }); + const controller = new AbortController(); + const pending = client.runMicrovm({ + imageIdentifier: 'arn:image', + maximumDurationSeconds: 120, + clientToken: 'exec-idempotent-1', + }, controller.signal); + controller.abort(new Error('job deadline')); + + await expect(pending).resolves.toMatchObject({ microvmId: 'mvm-recovered' }); + expect(sent).toHaveLength(2); + expect(sent.map(item => item.command.input.clientToken)) + .toEqual(['exec-idempotent-1', 'exec-idempotent-1']); + expect(sent[1].signal?.aborted).toBe(false); + }); + + test('does not reconcile a deterministic RunMicrovm rejection', async () => { + const { sender, sent } = stubSender([namedError('ValidationException')]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + await expect(client.runMicrovm({ + imageIdentifier: 'arn:image', + maximumDurationSeconds: 120, + clientToken: 'exec-invalid', + })).rejects.toMatchObject({ kind: 'validation', operation: 'RunMicrovm' }); + expect(sent).toHaveLength(1); + }); +}); + +describe('AwsLambdaMicrovmClient error classification', () => { + const cases: Array<[string, string]> = [ + ['ThrottlingException', 'throttled'], + ['TooManyRequestsException', 'throttled'], + ['ResourceNotFoundException', 'not_found'], + ['ConflictException', 'conflict'], + ['ServiceQuotaExceededException', 'quota_exceeded'], + ['ValidationException', 'validation'], + ['SomeUnknownException', 'other'], + ]; + + for (const [name, kind] of cases) { + test(`${name} -> ${kind}`, async () => { + const { sender } = stubSender([namedError(name)]); + const client = new AwsLambdaMicrovmClient({ client: sender }); + try { + await client.getMicrovm('mvm-1'); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(LambdaMicrovmApiError); + expect((error as LambdaMicrovmApiError).kind).toBe(kind as LambdaMicrovmApiError['kind']); + expect((error as LambdaMicrovmApiError).operation).toBe('GetMicrovm'); + } + }); + } +}); diff --git a/service/src/runtime-session/lambda-client-aws.ts b/service/src/runtime-session/lambda-client-aws.ts new file mode 100644 index 0000000..e2fb5b0 --- /dev/null +++ b/service/src/runtime-session/lambda-client-aws.ts @@ -0,0 +1,233 @@ +import { + LambdaMicrovmsClient, + RunMicrovmCommand, + GetMicrovmCommand, + SuspendMicrovmCommand, + ResumeMicrovmCommand, + TerminateMicrovmCommand, + CreateMicrovmAuthTokenCommand, + type MicrovmState, +} from '@aws-sdk/client-lambda-microvms'; +import { + LambdaMicrovmApiError, + MICROVM_AUTH_HEADER, + type LambdaMicrovmClient, + type LambdaMicrovmErrorKind, + type MicrovmAuthToken, + type MicrovmDescription, + type MicrovmLifecycleState, + type RunMicrovmArgs, +} from './lambda-client'; + +const THROTTLE_ERROR_NAMES = new Set(['ThrottlingException', 'TooManyRequestsException']); + +const ERROR_KIND_BY_NAME: Record = { + ResourceNotFoundException: 'not_found', + ConflictException: 'conflict', + ResourceConflictException: 'conflict', + ServiceQuotaExceededException: 'quota_exceeded', + ValidationException: 'validation', + InvalidParameterValueException: 'validation', +}; + +function classifyError(error: unknown): LambdaMicrovmErrorKind { + const name = (error as { name?: string } | null)?.name ?? ''; + if (THROTTLE_ERROR_NAMES.has(name)) return 'throttled'; + return ERROR_KIND_BY_NAME[name] ?? 'other'; +} + +function toDescription(response: { + microvmId?: string; + state?: MicrovmState; + endpoint?: string; + imageArn?: string; + imageVersion?: string; + maximumDurationInSeconds?: number; + startedAt?: Date; + stateReason?: string; +}): MicrovmDescription { + /* Every command (Run/Get/Suspend/Resume/Terminate) returns the VM id. A + * missing id means a partial/garbled response; fail fast rather than hand + * back `''`, which downstream getMicrovm('')/terminateMicrovm('') would act + * on — leaking the just-created VM as orphaned and billable. */ + if (response.microvmId == null || response.microvmId === '') { + throw new Error('Lambda MicroVM response omitted microvmId'); + } + return { + microvmId: response.microvmId, + state: (response.state ?? 'PENDING') as MicrovmLifecycleState, + endpoint: response.endpoint, + imageArn: response.imageArn, + imageVersion: response.imageVersion, + maximumDurationSeconds: response.maximumDurationInSeconds, + startedAtMs: response.startedAt?.getTime(), + stateReason: response.stateReason, + }; +} + +/** Minimal send-shaped surface so tests can stub the SDK client. */ +export interface MicrovmCommandSender { + send(command: unknown, options?: { abortSignal?: AbortSignal }): Promise; +} + +export class AwsLambdaMicrovmClient implements LambdaMicrovmClient { + private readonly client: MicrovmCommandSender; + private readonly requestTimeoutMs: number; + + constructor(options: { + region?: string; + client?: MicrovmCommandSender; + requestTimeoutMs?: number; + } = {}) { + this.client = options.client ?? new LambdaMicrovmsClient({ + region: options.region, + retryMode: 'adaptive', + maxAttempts: 3, + }); + this.requestTimeoutMs = Math.max(1, options.requestTimeoutMs ?? 60_000); + } + + private async send( + operation: string, + command: unknown, + callerSignal?: AbortSignal, + ): Promise { + /* Every SDK request has its own hard deadline even when the caller does not + * supply one (notably best-effort termination). When a job signal exists, + * either deadline cancels the adaptive-retry loop and underlying handler. */ + const timeoutSignal = AbortSignal.timeout(this.requestTimeoutMs); + const abortSignal = callerSignal + ? AbortSignal.any([callerSignal, timeoutSignal]) + : timeoutSignal; + try { + return await this.client.send(command, { abortSignal }) as T; + } catch (error) { + throw new LambdaMicrovmApiError( + classifyError(error), + operation, + (error as Error)?.message ?? `Lambda MicroVM ${operation} failed`, + error, + ); + } + } + + async runMicrovm(args: RunMicrovmArgs, signal?: AbortSignal): Promise { + const input = { + imageIdentifier: args.imageIdentifier, + imageVersion: args.imageVersion, + executionRoleArn: args.executionRoleArn, + ingressNetworkConnectors: args.ingressConnectorArns, + egressNetworkConnectors: args.egressConnectorArns, + maximumDurationInSeconds: args.maximumDurationSeconds, + idlePolicy: args.idlePolicy + ? { + maxIdleDurationSeconds: args.idlePolicy.maxIdleSeconds, + suspendedDurationSeconds: args.idlePolicy.suspendedSeconds, + autoResumeEnabled: args.idlePolicy.autoResume, + } + : undefined, + logging: args.logGroup ? { cloudWatch: { logGroup: args.logGroup } } : undefined, + runHookPayload: args.runHookPayload, + clientToken: args.clientToken, + }; + try { + const response = await this.send[0]>( + 'RunMicrovm', + new RunMicrovmCommand(input), + signal, + ); + return toDescription(response); + } catch (firstError) { + /* A timeout/abort or broken response is ambiguous: AWS may have accepted + * RunMicrovm before the client lost the response. Replay the SAME + * idempotency token once on an independent bounded request. If AWS + * accepted it, this recovers the MicroVM id so the backend can either + * continue or (when the job signal is already aborted) terminate it. + * Never reconcile deterministic validation/quota/throttle failures, and + * never replay a launch that lacks an idempotency token. */ + const ambiguous = + !(firstError instanceof LambdaMicrovmApiError) + || firstError.kind === 'other'; + if (!args.clientToken || !ambiguous) throw firstError; + try { + const recovered = await this.send[0]>( + 'RunMicrovmReconcile', + new RunMicrovmCommand(input), + ); + return toDescription(recovered); + } catch { + /* Preserve the original classification/message. The persisted session + * launch intent (stateful path) can make the same-token recovery on a + * successor; stateless VMs retain their short maximum duration. */ + throw firstError; + } + } + } + + async getMicrovm(microvmId: string, signal?: AbortSignal): Promise { + const response = await this.send[0]>( + 'GetMicrovm', + new GetMicrovmCommand({ microvmIdentifier: microvmId }), + signal, + ); + return toDescription(response); + } + + async suspendMicrovm(microvmId: string, signal?: AbortSignal): Promise { + await this.send( + 'SuspendMicrovm', + new SuspendMicrovmCommand({ microvmIdentifier: microvmId }), + signal, + ); + } + + async resumeMicrovm(microvmId: string, signal?: AbortSignal): Promise { + /* ResumeMicrovm's real response is empty. Read the resource afterward + * instead of passing that empty object through toDescription(), which + * would turn every successful resume into "response omitted microvmId". */ + await this.send( + 'ResumeMicrovm', + new ResumeMicrovmCommand({ microvmIdentifier: microvmId }), + signal, + ); + return this.getMicrovm(microvmId, signal); + } + + async terminateMicrovm(microvmId: string, signal?: AbortSignal): Promise { + await this.send( + 'TerminateMicrovm', + new TerminateMicrovmCommand({ microvmIdentifier: microvmId }), + signal, + ); + } + + async createMicrovmAuthToken(args: { + microvmId: string; + port: number; + ttlSeconds: number; + }, signal?: AbortSignal): Promise { + const expirationInMinutes = Math.min(Math.max(Math.ceil(args.ttlSeconds / 60), 1), 60); + const response = await this.send<{ authToken?: Record }>( + 'CreateMicrovmAuthToken', + new CreateMicrovmAuthTokenCommand({ + microvmIdentifier: args.microvmId, + expirationInMinutes, + allowedPorts: [{ port: args.port }], + }), + signal, + ); + const token = response.authToken?.[MICROVM_AUTH_HEADER]; + if (token == null || token.length === 0) { + throw new LambdaMicrovmApiError( + 'other', + 'CreateMicrovmAuthToken', + `CreateMicrovmAuthToken response missing ${MICROVM_AUTH_HEADER} entry`, + ); + } + return { + headerName: MICROVM_AUTH_HEADER, + token, + expiresAtMs: Date.now() + expirationInMinutes * 60_000, + }; + } +} diff --git a/service/src/runtime-session/lambda-client-fake.test.ts b/service/src/runtime-session/lambda-client-fake.test.ts new file mode 100644 index 0000000..fa7b24e --- /dev/null +++ b/service/src/runtime-session/lambda-client-fake.test.ts @@ -0,0 +1,76 @@ +import { describe, expect, test } from 'bun:test'; +import { FakeLambdaMicrovmClient } from './lambda-client-fake'; +import { LambdaMicrovmApiError } from './lambda-client'; + +const RUN_ARGS = { + imageIdentifier: 'arn:image/codeapi', + maximumDurationSeconds: 28_800, +}; + +describe('FakeLambdaMicrovmClient state machine', () => { + test('launches RUNNING by default with a per-VM endpoint', async () => { + const fake = new FakeLambdaMicrovmClient({ endpointProvider: (id) => `http://localhost:9/${id}` }); + const vm = await fake.runMicrovm(RUN_ARGS); + expect(vm.state).toBe('RUNNING'); + expect(vm.endpoint).toBe(`http://localhost:9/${vm.microvmId}`); + expect((await fake.getMicrovm(vm.microvmId)).state).toBe('RUNNING'); + }); + + test('delayNextLaunch keeps the VM PENDING for N polls', async () => { + const fake = new FakeLambdaMicrovmClient(); + fake.delayNextLaunch(2); + const vm = await fake.runMicrovm(RUN_ARGS); + expect(vm.state).toBe('PENDING'); + expect((await fake.getMicrovm(vm.microvmId)).state).toBe('PENDING'); + expect((await fake.getMicrovm(vm.microvmId)).state).toBe('RUNNING'); + }); + + test('suspend/resume/terminate transitions', async () => { + const fake = new FakeLambdaMicrovmClient(); + const vm = await fake.runMicrovm(RUN_ARGS); + + await fake.suspendMicrovm(vm.microvmId); + expect((await fake.getMicrovm(vm.microvmId)).state).toBe('SUSPENDED'); + + expect((await fake.resumeMicrovm(vm.microvmId)).state).toBe('RUNNING'); + + await fake.terminateMicrovm(vm.microvmId); + expect((await fake.getMicrovm(vm.microvmId)).state).toBe('TERMINATED'); + expect(fake.resumeMicrovm(vm.microvmId)).rejects.toThrow('is TERMINATED'); + }); + + test('clientToken makes launch idempotent', async () => { + const fake = new FakeLambdaMicrovmClient(); + const first = await fake.runMicrovm({ ...RUN_ARGS, clientToken: 'launch-1' }); + const second = await fake.runMicrovm({ ...RUN_ARGS, clientToken: 'launch-1' }); + const third = await fake.runMicrovm({ ...RUN_ARGS, clientToken: 'launch-2' }); + expect(second.microvmId).toBe(first.microvmId); + expect(third.microvmId).not.toBe(first.microvmId); + expect(fake.vms.size).toBe(2); + }); + + test('failNext raises once then recovers, and unknown VMs are not_found', async () => { + const fake = new FakeLambdaMicrovmClient(); + fake.failNext('runMicrovm'); + expect(fake.runMicrovm(RUN_ARGS)).rejects.toThrow('scripted'); + const vm = await fake.runMicrovm(RUN_ARGS); + expect(vm.state).toBe('RUNNING'); + + try { + await fake.getMicrovm('missing'); + throw new Error('expected rejection'); + } catch (error) { + expect((error as LambdaMicrovmApiError).kind).toBe('not_found'); + } + }); + + test('mints distinct tokens per VM and records calls', async () => { + const fake = new FakeLambdaMicrovmClient(); + const vm = await fake.runMicrovm(RUN_ARGS); + const first = await fake.createMicrovmAuthToken({ microvmId: vm.microvmId, port: 8080, ttlSeconds: 300 }); + const second = await fake.createMicrovmAuthToken({ microvmId: vm.microvmId, port: 8080, ttlSeconds: 300 }); + expect(first.token).not.toBe(second.token); + expect(fake.vms.get(vm.microvmId)?.mintedTokens).toEqual([first.token, second.token]); + expect(fake.callsFor('createMicrovmAuthToken')).toHaveLength(2); + }); +}); diff --git a/service/src/runtime-session/lambda-client-fake.ts b/service/src/runtime-session/lambda-client-fake.ts new file mode 100644 index 0000000..f9af670 --- /dev/null +++ b/service/src/runtime-session/lambda-client-fake.ts @@ -0,0 +1,203 @@ +import { nanoid } from 'nanoid'; +import { + LambdaMicrovmApiError, + MICROVM_AUTH_HEADER, + type LambdaMicrovmClient, + type MicrovmAuthToken, + type MicrovmDescription, + type MicrovmIdlePolicy, + type MicrovmLifecycleState, + type RunMicrovmArgs, +} from './lambda-client'; + +export interface FakeMicrovm { + microvmId: string; + state: MicrovmLifecycleState; + endpoint: string; + imageIdentifier: string; + imageVersion?: string; + maximumDurationSeconds: number; + idlePolicy?: MicrovmIdlePolicy; + runHookPayload?: string; + clientToken?: string; + startedAtMs: number; + mintedTokens: string[]; +} + +type FakeOp = 'runMicrovm' | 'getMicrovm' | 'suspendMicrovm' | 'resumeMicrovm' | 'terminateMicrovm' | 'createMicrovmAuthToken'; + +/** + * In-memory control-plane fake for bun tests. Transport-free: the test + * supplies `endpointProvider` (usually a Bun.serve URL) so the backend's real + * HTTP proxy path is exercised against a fake sandbox endpoint. + * + * Launch behavior: VMs come up RUNNING immediately unless + * `launchStates` supplies an explicit state sequence (e.g. keep a VM + * PENDING for N getMicrovm polls). + */ +export class FakeLambdaMicrovmClient implements LambdaMicrovmClient { + readonly vms = new Map(); + readonly calls: Array<{ op: FakeOp; args: unknown }> = []; + private readonly failures = new Map(); + private pendingPollsByClientToken = new Map(); + private vmSeq = 0; + + constructor( + private readonly options: { + endpointProvider?: (microvmId: string) => string; + nowFn?: () => number; + } = {}, + ) {} + + /** Queue an error for the next call of `op` (FIFO). */ + failNext(op: FakeOp, error?: Error): void { + const queue = this.failures.get(op) ?? []; + queue.push(error ?? new LambdaMicrovmApiError('other', op, `${op} failed (scripted)`)); + this.failures.set(op, queue); + } + + /** Make the next launched VM stay PENDING for `polls` getMicrovm calls. */ + delayNextLaunch(polls: number): void { + this.pendingPollsByClientToken.set('__next__', polls); + } + + /** Make the next launched VM come back already TERMINATED (boot-time death). */ + terminateNextLaunch(): void { + this.terminateNextLaunches += 1; + } + + private terminateNextLaunches = 0; + + setState(microvmId: string, state: MicrovmLifecycleState): void { + const vm = this.mustGet(microvmId); + vm.state = state; + } + + private now(): number { + return this.options.nowFn?.() ?? Date.now(); + } + + private takeFailure(op: FakeOp): void { + const queue = this.failures.get(op); + const error = queue?.shift(); + if (error) throw error; + } + + private mustGet(microvmId: string): FakeMicrovm { + const vm = this.vms.get(microvmId); + if (!vm) { + throw new LambdaMicrovmApiError('not_found', 'GetMicrovm', `MicroVM ${microvmId} not found`); + } + return vm; + } + + private describe(vm: FakeMicrovm): MicrovmDescription { + return { + microvmId: vm.microvmId, + state: vm.state, + endpoint: vm.endpoint, + imageArn: vm.imageIdentifier, + imageVersion: vm.imageVersion, + maximumDurationSeconds: vm.maximumDurationSeconds, + startedAtMs: vm.startedAtMs, + }; + } + + async runMicrovm(args: RunMicrovmArgs): Promise { + this.calls.push({ op: 'runMicrovm', args }); + this.takeFailure('runMicrovm'); + + if (args.clientToken != null) { + const existing = [...this.vms.values()].find((vm) => vm.clientToken === args.clientToken); + if (existing) return this.describe(existing); + } + + const microvmId = `fake-mvm-${++this.vmSeq}-${nanoid(6)}`; + const pendingPolls = this.pendingPollsByClientToken.get('__next__') ?? 0; + this.pendingPollsByClientToken.delete('__next__'); + if (pendingPolls > 0) { + this.pendingPollsByClientToken.set(microvmId, pendingPolls); + } + const bootDeath = this.terminateNextLaunches > 0; + if (bootDeath) { + this.terminateNextLaunches -= 1; + } + + const vm: FakeMicrovm = { + microvmId, + state: bootDeath ? 'TERMINATED' : pendingPolls > 0 ? 'PENDING' : 'RUNNING', + endpoint: this.options.endpointProvider?.(microvmId) ?? `https://${microvmId}.fake-microvm.on.aws`, + imageIdentifier: args.imageIdentifier, + imageVersion: args.imageVersion, + maximumDurationSeconds: args.maximumDurationSeconds, + idlePolicy: args.idlePolicy, + runHookPayload: args.runHookPayload, + clientToken: args.clientToken, + startedAtMs: this.now(), + mintedTokens: [], + }; + this.vms.set(microvmId, vm); + return this.describe(vm); + } + + async getMicrovm(microvmId: string): Promise { + this.calls.push({ op: 'getMicrovm', args: { microvmId } }); + this.takeFailure('getMicrovm'); + const vm = this.mustGet(microvmId); + const remaining = this.pendingPollsByClientToken.get(microvmId); + if (remaining != null) { + if (remaining <= 1) { + this.pendingPollsByClientToken.delete(microvmId); + vm.state = 'RUNNING'; + } else { + this.pendingPollsByClientToken.set(microvmId, remaining - 1); + } + } + return this.describe(vm); + } + + async suspendMicrovm(microvmId: string): Promise { + this.calls.push({ op: 'suspendMicrovm', args: { microvmId } }); + this.takeFailure('suspendMicrovm'); + this.mustGet(microvmId).state = 'SUSPENDED'; + } + + async resumeMicrovm(microvmId: string): Promise { + this.calls.push({ op: 'resumeMicrovm', args: { microvmId } }); + this.takeFailure('resumeMicrovm'); + const vm = this.mustGet(microvmId); + if (vm.state === 'TERMINATED' || vm.state === 'TERMINATING') { + throw new LambdaMicrovmApiError('conflict', 'ResumeMicrovm', `MicroVM ${microvmId} is ${vm.state}`); + } + vm.state = 'RUNNING'; + return this.describe(vm); + } + + async terminateMicrovm(microvmId: string): Promise { + this.calls.push({ op: 'terminateMicrovm', args: { microvmId } }); + this.takeFailure('terminateMicrovm'); + const vm = this.vms.get(microvmId); + if (vm) vm.state = 'TERMINATED'; + } + + async createMicrovmAuthToken(args: { + microvmId: string; + port: number; + ttlSeconds: number; + }): Promise { + this.calls.push({ op: 'createMicrovmAuthToken', args }); + this.takeFailure('createMicrovmAuthToken'); + const vm = this.mustGet(args.microvmId); + const token = `fake-proxy-token-${args.microvmId}-${vm.mintedTokens.length + 1}`; + vm.mintedTokens.push(token); + return { + headerName: MICROVM_AUTH_HEADER, + token, + expiresAtMs: this.now() + args.ttlSeconds * 1_000, + }; + } + + callsFor(op: FakeOp): Array<{ op: FakeOp; args: unknown }> { + return this.calls.filter((call) => call.op === op); + } +} diff --git a/service/src/runtime-session/lambda-client.ts b/service/src/runtime-session/lambda-client.ts new file mode 100644 index 0000000..48d7de0 --- /dev/null +++ b/service/src/runtime-session/lambda-client.ts @@ -0,0 +1,99 @@ +/** + * Thin, fakeable wrapper over the AWS Lambda MicroVM control plane. + * `lambda-client-aws.ts` is the ONLY module allowed to import `@aws-sdk/*`; + * everything else (backend, sweeper, tests) programs against this interface. + */ + +export type MicrovmLifecycleState = + | 'PENDING' + | 'RUNNING' + | 'SUSPENDING' + | 'SUSPENDED' + | 'TERMINATING' + | 'TERMINATED'; + +export interface MicrovmDescription { + microvmId: string; + state: MicrovmLifecycleState; + endpoint?: string; + imageArn?: string; + imageVersion?: string; + maximumDurationSeconds?: number; + startedAtMs?: number; + stateReason?: string; +} + +export interface MicrovmIdlePolicy { + maxIdleSeconds: number; + suspendedSeconds: number; + autoResume: boolean; +} + +export interface RunMicrovmArgs { + imageIdentifier: string; + imageVersion?: string; + executionRoleArn?: string; + ingressConnectorArns?: string[]; + egressConnectorArns?: string[]; + maximumDurationSeconds: number; + idlePolicy?: MicrovmIdlePolicy; + /** CloudWatch log group for the VM's stdout/stderr. Needs an executionRoleArn + * too, or the logs go nowhere. */ + logGroup?: string; + /** Delivered verbatim as the /run lifecycle hook body (AWS cap: 16KB). */ + runHookPayload?: string; + /** Idempotency token so a retried launch cannot double-provision. */ + clientToken?: string; +} + +export interface MicrovmAuthToken { + /** Header name the MicroVM proxy expects; AWS returns a map keyed by it. */ + headerName: string; + token: string; + expiresAtMs: number; +} + +export type LambdaMicrovmErrorKind = + | 'throttled' + | 'not_found' + | 'conflict' + | 'quota_exceeded' + | 'validation' + | 'other'; + +export class LambdaMicrovmApiError extends Error { + constructor( + public readonly kind: LambdaMicrovmErrorKind, + public readonly operation: string, + message: string, + public readonly cause?: unknown, + ) { + super(message); + this.name = 'LambdaMicrovmApiError'; + } +} + +export const MICROVM_AUTH_HEADER = 'X-aws-proxy-auth'; + +/** MicroVM endpoint traffic defaults to port 8080; to reach a different target + * port the request must carry this header (AWS routes to 8080 without it). */ +export const MICROVM_PORT_HEADER = 'X-aws-proxy-port'; +export const DEFAULT_MICROVM_PORT = 8080; + +/** The port-routing header, only when the target port isn't the 8080 default. */ +export function microvmPortHeaders(port: number): Record { + return port === DEFAULT_MICROVM_PORT ? {} : { [MICROVM_PORT_HEADER]: String(port) }; +} + +export interface LambdaMicrovmClient { + runMicrovm(args: RunMicrovmArgs, signal?: AbortSignal): Promise; + getMicrovm(microvmId: string, signal?: AbortSignal): Promise; + suspendMicrovm(microvmId: string, signal?: AbortSignal): Promise; + resumeMicrovm(microvmId: string, signal?: AbortSignal): Promise; + terminateMicrovm(microvmId: string, signal?: AbortSignal): Promise; + createMicrovmAuthToken(args: { + microvmId: string; + port: number; + ttlSeconds: number; + }, signal?: AbortSignal): Promise; +} diff --git a/service/src/runtime-session/lock-heartbeat.test.ts b/service/src/runtime-session/lock-heartbeat.test.ts new file mode 100644 index 0000000..83d7137 --- /dev/null +++ b/service/src/runtime-session/lock-heartbeat.test.ts @@ -0,0 +1,90 @@ +import { describe, expect, test } from 'bun:test'; +import { startRuntimeSessionLockHeartbeat } from './lock-heartbeat'; + +const wait = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); + +describe('runtime session lock heartbeat', () => { + test('independently fences a holder whose renewal never settles', async () => { + const fence = new AbortController(); + let calls = 0; + const heartbeat = startRuntimeSessionLockHeartbeat({ + renew: () => { + calls += 1; + return new Promise(() => {}); + }, + fence, + ttlMs: 40, + intervalMs: 5, + }); + try { + await wait(70); + expect(fence.signal.aborted).toBe(true); + expect(calls).toBe(1); + } finally { + heartbeat.stop(); + } + }); + + test('fences immediately when Redis reports the lock was lost', async () => { + const fence = new AbortController(); + const heartbeat = startRuntimeSessionLockHeartbeat({ + renew: async () => 'lost', + fence, + ttlMs: 1_000, + intervalMs: 5, + }); + try { + await wait(20); + expect(fence.signal.aborted).toBe(true); + } finally { + heartbeat.stop(); + } + }); + + test('an unexpected renewal rejection still fences at the independent deadline', async () => { + const fence = new AbortController(); + const heartbeat = startRuntimeSessionLockHeartbeat({ + renew: async () => { + throw new Error('transport exploded'); + }, + fence, + ttlMs: 40, + intervalMs: 5, + }); + try { + await wait(70); + expect(fence.signal.aborted).toBe(true); + } finally { + heartbeat.stop(); + } + }); + + test('a renewal that settles after stop cannot re-arm the expiry watchdog', async () => { + const fence = new AbortController(); + let resolveRenewal!: (result: 'held') => void; + let markRenewalStarted!: () => void; + const renewalStarted = new Promise((resolve) => { + markRenewalStarted = resolve; + }); + const heartbeat = startRuntimeSessionLockHeartbeat({ + renew: () => { + markRenewalStarted(); + return new Promise<'held'>((resolve) => { + resolveRenewal = resolve; + }); + }, + fence, + ttlMs: 40, + intervalMs: 5, + }); + try { + await renewalStarted; + heartbeat.stop(); + resolveRenewal('held'); + await wait(70); + expect(fence.signal.aborted).toBe(false); + } finally { + heartbeat.stop(); + } + }); +}); diff --git a/service/src/runtime-session/lock-heartbeat.ts b/service/src/runtime-session/lock-heartbeat.ts new file mode 100644 index 0000000..b10d9d7 --- /dev/null +++ b/service/src/runtime-session/lock-heartbeat.ts @@ -0,0 +1,61 @@ +export interface LockHeartbeat { + stop(): void; +} + +/** + * Renews a session lease without overlapping Redis calls and independently + * fences the holder when no renewal can be confirmed before the current TTL. + * The deadline is measured from request start, so a delayed "held" response + * cannot extend a lease that may already have expired in the meantime. + */ +export function startRuntimeSessionLockHeartbeat(args: { + renew: () => Promise<'held' | 'lost' | 'error'>; + fence: AbortController; + ttlMs: number; + intervalMs?: number; +}): LockHeartbeat { + const ttlMs = args.ttlMs; + const intervalMs = args.intervalMs ?? Math.floor(ttlMs / 3); + let stopped = false; + let inFlight = false; + let expiry: ReturnType; + + const armExpiry = (remainingMs = ttlMs): void => { + clearTimeout(expiry); + expiry = setTimeout(() => args.fence.abort(), Math.max(0, remainingMs)); + }; + armExpiry(); + + const tick = (): void => { + if (stopped || inFlight || args.fence.signal.aborted) return; + inFlight = true; + const startedAt = Date.now(); + void args.renew().then( + (renewal) => { + if (stopped || args.fence.signal.aborted) return; + if (renewal === 'lost') { + args.fence.abort(); + return; + } + if (renewal === 'held') { + armExpiry(ttlMs - (Date.now() - startedAt)); + } + }, + () => { + /* Treat an unexpected rejection like the registry's explicit `error` + * result: the independent TTL watchdog remains authoritative. */ + }, + ).finally(() => { + inFlight = false; + }); + }; + + const interval = setInterval(tick, intervalMs); + return { + stop(): void { + stopped = true; + clearInterval(interval); + clearTimeout(expiry); + }, + }; +} diff --git a/service/src/runtime-session/registry.test.ts b/service/src/runtime-session/registry.test.ts new file mode 100644 index 0000000..97e35c9 --- /dev/null +++ b/service/src/runtime-session/registry.test.ts @@ -0,0 +1,174 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import { + acquireRuntimeSessionLock, + allocateCheckpointSequence, + allocateRuntimeSessionGeneration, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + removeRuntimeSession, + resetRedisForTests, + setRedisForTests, + waitForRuntimeSessionLock, + writeRuntimeSessionRecord, + type RuntimeSessionRecord, +} from './registry'; + +let mock: InstanceType; + +beforeEach(async () => { + /* ioredis-mock shares one keyspace across instances — flush per test. */ + mock = new RedisMock(); + await mock.flushall(); + setRedisForTests(mock); +}); + +afterEach(() => { + resetRedisForTests(); +}); + +function record(overrides: Partial = {}): RuntimeSessionRecord { + return { + runtime_session_id: 'rt_abc123', + tenant_id: 'tenant-a', + canonical_user_id: 'user-1', + state: 'PENDING', + generation: 1, + last_seen_at: 1_778_250_000_000, + ...overrides, + }; +} + +describe('runtime session lock', () => { + test('acquire is exclusive; release makes it available again', async () => { + const token = await acquireRuntimeSessionLock('rt_abc123'); + expect(token).not.toBeNull(); + expect(await acquireRuntimeSessionLock('rt_abc123')).toBeNull(); + await releaseRuntimeSessionLock('rt_abc123', token as string); + expect(await acquireRuntimeSessionLock('rt_abc123')).not.toBeNull(); + }); + + test('release is CAS-guarded: a stale token cannot free the current holder', async () => { + const first = await acquireRuntimeSessionLock('rt_abc123'); + await releaseRuntimeSessionLock('rt_abc123', first as string); + const second = await acquireRuntimeSessionLock('rt_abc123'); + await releaseRuntimeSessionLock('rt_abc123', first as string); + expect(await acquireRuntimeSessionLock('rt_abc123')).toBeNull(); + await releaseRuntimeSessionLock('rt_abc123', second as string); + }); + + test('waitForRuntimeSessionLock polls until the holder releases', async () => { + const holder = await acquireRuntimeSessionLock('rt_abc123'); + setTimeout(() => void releaseRuntimeSessionLock('rt_abc123', holder as string), 60); + const token = await waitForRuntimeSessionLock('rt_abc123', { waitMs: 2_000, pollMs: 20 }); + expect(token).not.toBeNull(); + }); + + test('waitForRuntimeSessionLock gives up after waitMs', async () => { + await acquireRuntimeSessionLock('rt_abc123'); + const started = Date.now(); + const token = await waitForRuntimeSessionLock('rt_abc123', { waitMs: 120, pollMs: 25 }); + expect(token).toBeNull(); + expect(Date.now() - started).toBeLessThan(1_000); + }); + + test('waitForRuntimeSessionLock stops promptly when the job is canceled', async () => { + const holder = await acquireRuntimeSessionLock('rt_abc123'); + const controller = new AbortController(); + const started = Date.now(); + setTimeout(() => controller.abort(new Error('job deadline')), 10); + + await expect(waitForRuntimeSessionLock('rt_abc123', { + waitMs: 5_000, + pollMs: 1_000, + signal: controller.signal, + })).rejects.toThrow('job deadline'); + expect(Date.now() - started).toBeLessThan(1_000); + expect(await acquireRuntimeSessionLock('rt_abc123')).toBeNull(); + await releaseRuntimeSessionLock('rt_abc123', holder as string); + }); +}); + +describe('fenced record writes', () => { + test('retries a transient lock-release failure so the session is immediately reusable', async () => { + const token = (await acquireRuntimeSessionLock('rt_release_retry')) as string; + const scripted = mock as unknown as { + releaseRuntimeSessionLockScript( + lockKey: string, + lockToken: string, + ): Promise; + }; + const release = scripted.releaseRuntimeSessionLockScript.bind(scripted); + let calls = 0; + scripted.releaseRuntimeSessionLockScript = async (lockKey, lockToken) => { + calls += 1; + if (calls === 1) throw new Error('temporary Redis failover'); + return release(lockKey, lockToken); + }; + + await releaseRuntimeSessionLock('rt_release_retry', token); + + expect(calls).toBe(2); + expect(await acquireRuntimeSessionLock('rt_release_retry')).not.toBeNull(); + }); + + test('write succeeds while holding the lock and round-trips the record', async () => { + const token = (await acquireRuntimeSessionLock('rt_abc123')) as string; + const rec = record({ state: 'RUNNING', microvm_id: 'mvm-1', endpoint: 'https://vm.example', generation: 3 }); + expect(await writeRuntimeSessionRecord(rec, token)).toBe(true); + expect(await readRuntimeSessionRecord('rt_abc123')).toEqual(rec); + }); + + test('reads a corrupt record as missing instead of throwing', async () => { + await mock.set('rtsx:sess:rt_bad', '{not valid json'); + expect(await readRuntimeSessionRecord('rt_bad')).toBeNull(); + }); + + test('write is fenced after the lock is lost', async () => { + const token = (await acquireRuntimeSessionLock('rt_abc123')) as string; + await releaseRuntimeSessionLock('rt_abc123', token); + const thief = await acquireRuntimeSessionLock('rt_abc123'); + expect(thief).not.toBeNull(); + expect(await writeRuntimeSessionRecord(record(), token)).toBe(false); + expect(await readRuntimeSessionRecord('rt_abc123')).toBeNull(); + }); + + test('write is fenced when no lock exists at all', async () => { + expect(await writeRuntimeSessionRecord(record(), 'never-held')).toBe(false); + }); + + test('removal is fenced and clears the record', async () => { + const token = (await acquireRuntimeSessionLock('rt_abc123')) as string; + await writeRuntimeSessionRecord(record(), token); + + expect(await removeRuntimeSession('rt_abc123', 'stale-token')).toBe(false); + expect(await readRuntimeSessionRecord('rt_abc123')).not.toBeNull(); + + expect(await removeRuntimeSession('rt_abc123', token)).toBe(true); + expect(await readRuntimeSessionRecord('rt_abc123')).toBeNull(); + }); +}); + +describe('generation counter', () => { + test('increments monotonically per session and independently across sessions', async () => { + expect(await allocateRuntimeSessionGeneration('rt_abc123')).toBe(1); + expect(await allocateRuntimeSessionGeneration('rt_abc123')).toBe(2); + expect(await allocateRuntimeSessionGeneration('rt_abc123')).toBe(3); + expect(await allocateRuntimeSessionGeneration('rt_other')).toBe(1); + }); +}); + +describe('checkpoint sequence counter', () => { + test('concurrent stale holders reserve distinct keys above the durable high-water mark', async () => { + /* Models two holders that both listed durable sequence 100 around a lease + * handoff. A split INCR + reseed SET can return 101 to both and let the + * stale upload overwrite the new holder's committed object. The atomic + * reservation must serialize them as 101 and 102 instead. */ + const reservations = await Promise.all([ + allocateCheckpointSequence('rt_abc123', 100), + allocateCheckpointSequence('rt_abc123', 100), + ]); + expect(reservations.sort((a, b) => a - b)).toEqual([101, 102]); + expect(await allocateCheckpointSequence('rt_abc123', 50)).toBe(103); + }); +}); diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts new file mode 100644 index 0000000..da2ba2e --- /dev/null +++ b/service/src/runtime-session/registry.ts @@ -0,0 +1,310 @@ +import { nanoid } from 'nanoid'; +import type { Redis } from 'ioredis'; +import { connection } from '../queue'; +import { env } from '../config'; +import logger from '../logger'; + +/** + * Redis-backed registry mapping a `runtime_session_id` to its live (or + * suspended) Lambda MicroVM. Keys: + * + * rtsx:sess: JSON RuntimeSessionRecord (TTL: record TTL) + * rtsx:lock: per-session mutex token (SET NX PX) + * rtsx:gen: monotonic generation counter (INCR) (TTL: record TTL) + * + * Fencing: every record mutation runs through a Lua script that checks the + * caller still holds the session lock. A `false` return means the caller was + * fenced (lock expired or stolen) and must treat any MicroVM it just launched + * as an orphan to terminate. Lua stays within the GET/SET/DEL string-compare + * subset that ioredis-mock supports (see replay-state.ts). + */ + +export type RuntimeSessionState = 'PENDING' | 'RUNNING' | 'SUSPENDED' | 'TERMINATING' | 'TERMINATED'; + +export interface RuntimeSessionRecord { + runtime_session_id: string; + tenant_id: string; + canonical_user_id: string; + microvm_id?: string; + endpoint?: string; + port?: number; + image_arn?: string; + image_version?: string; + /** Fingerprint of every immutable launch/security input. A deploy that + * changes one must not reuse a VM launched under the old policy. */ + launch_fingerprint?: string; + state: RuntimeSessionState; + generation: number; + launched_at?: number; + last_seen_at: number; + hard_deadline_at?: number; + workspace_checkpoint?: string; + checkpointed_at?: number; + last_error?: string; +} + +const SESS_PREFIX = 'rtsx:sess:'; +const LOCK_PREFIX = 'rtsx:lock:'; +const GEN_PREFIX = 'rtsx:gen:'; +const CKPT_SEQ_PREFIX = 'rtsx:ckptseq:'; + +/** The session lock is held across the WHOLE `executeSession` critical path + * (launch throttle, readiness/restore, execute, post-run checkpoint), which sums + * to a large and variable worst case once per-op token-mint throttle waits are + * included. Rather than pin the TTL to that sum (fragile — a missed term lets a + * second worker fence a live holder and mutate the session concurrently), the + * holder RENEWS the lock on a heartbeat (`renewRuntimeSessionLock`) for as long + * as it runs. This value is therefore just a comfortable BASE that must outlive + * one heartbeat interval plus a stalled event loop — it already covers a normal + * relaunch (execute + launch + health + the checkpoint I/Os) with headroom. */ +export const RUNTIME_SESSION_LOCK_TTL_MS = + env.JOB_TIMEOUT + + 2 * env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + + env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS + + 4 * env.CHECKPOINT_TIMEOUT_MS + + 60_000; + +const MAX_MICROVM_DURATION_SECONDS = 28_800; +export const RUNTIME_SESSION_RECORD_TTL_SECONDS = MAX_MICROVM_DURATION_SECONDS + 600; + +type RedisWithScripts = Redis & { + releaseRuntimeSessionLockScript(lockKey: string, token: string): Promise; + renewRuntimeSessionLockScript(lockKey: string, token: string, ttlMs: string): Promise; + writeRuntimeSessionRecordScript( + sessKey: string, + lockKey: string, + token: string, + recordJson: string, + ttlSeconds: string, + ): Promise; + removeRuntimeSessionScript( + sessKey: string, + lockKey: string, + token: string, + ): Promise; + allocateCheckpointSequenceScript( + sequenceKey: string, + retainedMax: string, + ttlSeconds: string, + ): Promise; +}; + +const SCRIPTS_REGISTERED = Symbol.for('runtime-session-registry.scriptsRegistered'); + +function registerScripts(client: Redis): RedisWithScripts { + const tagged = client as Redis & { [SCRIPTS_REGISTERED]?: true }; + if (tagged[SCRIPTS_REGISTERED]) return client as RedisWithScripts; + client.defineCommand('releaseRuntimeSessionLockScript', { + numberOfKeys: 1, + lua: "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", + }); + client.defineCommand('renewRuntimeSessionLockScript', { + numberOfKeys: 1, + lua: "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('pexpire', KEYS[1], ARGV[2]) else return 0 end", + }); + client.defineCommand('writeRuntimeSessionRecordScript', { + numberOfKeys: 2, + lua: `if redis.call('get', KEYS[2]) == ARGV[1] then + redis.call('set', KEYS[1], ARGV[2], 'EX', ARGV[3]) + return 1 +else + return 0 +end`, + }); + client.defineCommand('removeRuntimeSessionScript', { + numberOfKeys: 2, + lua: `if redis.call('get', KEYS[2]) == ARGV[1] then + redis.call('del', KEYS[1]) + return 1 +else + return 0 +end`, + }); + client.defineCommand('allocateCheckpointSequenceScript', { + numberOfKeys: 1, + lua: `local current = tonumber(redis.call('get', KEYS[1]) or '0') +local retained = tonumber(ARGV[1]) +if retained > current then current = retained end +local sequence = current + 1 +redis.call('set', KEYS[1], tostring(sequence), 'EX', ARGV[2]) +return sequence`, + }); + tagged[SCRIPTS_REGISTERED] = true; + return client as RedisWithScripts; +} + +let redis: RedisWithScripts = registerScripts(connection); + +/** Test seam mirroring replay-state.ts: swap in ioredis-mock per test. */ +export function setRedisForTests(client: Redis): void { + redis = registerScripts(client); +} + +export function resetRedisForTests(): void { + redis = registerScripts(connection); +} + +export async function acquireRuntimeSessionLock( + runtimeSessionId: string, + ttlMs: number = RUNTIME_SESSION_LOCK_TTL_MS, +): Promise { + const token = nanoid(); + const result = await redis.set(`${LOCK_PREFIX}${runtimeSessionId}`, token, 'PX', ttlMs, 'NX'); + return result === 'OK' ? token : null; +} + +/** Polls for the session mutex; returns null once `waitMs` is exhausted. + * Stateful callers surface contention as HTTP 409 instead of executing cold. */ +export async function waitForRuntimeSessionLock( + runtimeSessionId: string, + args: { waitMs: number; pollMs?: number; ttlMs?: number; signal?: AbortSignal }, +): Promise { + const pollMs = args.pollMs ?? 250; + const deadline = Date.now() + args.waitMs; + for (;;) { + args.signal?.throwIfAborted(); + const token = await acquireRuntimeSessionLock(runtimeSessionId, args.ttlMs); + if (args.signal?.aborted) { + if (token != null) await releaseRuntimeSessionLock(runtimeSessionId, token); + args.signal.throwIfAborted(); + } + if (token != null) return token; + if (Date.now() + pollMs > deadline) return null; + await new Promise((resolve, reject) => { + const timer = setTimeout(() => { + args.signal?.removeEventListener('abort', onAbort); + resolve(); + }, pollMs); + const onAbort = (): void => { + clearTimeout(timer); + reject(args.signal?.reason instanceof Error + ? args.signal.reason + : new Error('Runtime session lock wait aborted')); + }; + args.signal?.addEventListener('abort', onAbort, { once: true }); + }); + } +} + +export async function releaseRuntimeSessionLock(runtimeSessionId: string, token: string): Promise { + const retryDelaysMs = [50, 200] as const; + let lastError: unknown; + for (let attempt = 0; attempt <= retryDelaysMs.length; attempt += 1) { + try { + await redis.releaseRuntimeSessionLockScript(`${LOCK_PREFIX}${runtimeSessionId}`, token); + return; + } catch (error) { + lastError = error; + if (attempt < retryDelaysMs.length) { + await new Promise(resolve => setTimeout(resolve, retryDelaysMs[attempt])); + } + } + } + /* Release runs from finally blocks. Do not mask either a successful execute + * or its primary error, but heal brief Redis failovers before accepting that + * the token-guarded lease must age out. The Lua release is idempotent, so a + * retry is safe even if Redis deleted the key but lost the first response. */ + logger.warn('Failed to release runtime session lock after retries', { + runtimeSessionId, + err: lastError, + }); +} + +/** Fenced heartbeat: extend the lock's TTL only while we still hold the token. + * Returns false if we've been fenced (another worker owns the lock now), which + * the caller uses to stop renewing. Lets the critical path run arbitrarily long + * (launch throttle + restore + execute + checkpoint, each with its own I/O and + * token-mint waits) without the TTL having to bound the worst-case sum. */ +/** `lost` is positive evidence another holder fenced us (token mismatch); + * `error` is a transport failure where the lock may well still be held — + * callers must only abort in-flight work on `lost`, never on a single + * transient `error` (the TTL is a multiple of the heartbeat interval, so + * the next tick retries well before expiry). */ +export type LockRenewal = 'held' | 'lost' | 'error'; + +export async function renewRuntimeSessionLock( + runtimeSessionId: string, + token: string, + ttlMs: number = RUNTIME_SESSION_LOCK_TTL_MS, +): Promise { + try { + const result = await redis.renewRuntimeSessionLockScript( + `${LOCK_PREFIX}${runtimeSessionId}`, + token, + String(ttlMs), + ); + return result === 1 ? 'held' : 'lost'; + } catch (err) { + logger.warn('Failed to renew runtime session lock', { runtimeSessionId, err }); + return 'error'; + } +} + +export async function readRuntimeSessionRecord(runtimeSessionId: string): Promise { + const data = await redis.get(`${SESS_PREFIX}${runtimeSessionId}`); + if (data == null) return null; + /* Treat a corrupt/incompatible record as missing so a single bad key can't + * wedge every request for the session until it is manually deleted. */ + try { + return JSON.parse(data) as RuntimeSessionRecord; + } catch (err) { + logger.warn('Discarding malformed runtime session record', { runtimeSessionId, err }); + return null; + } +} + +/** Fenced write: persists the record only while `lockToken` still holds the + * session mutex. Returns false when the caller was fenced. */ +export async function writeRuntimeSessionRecord( + record: RuntimeSessionRecord, + lockToken: string, + ttlSeconds: number = RUNTIME_SESSION_RECORD_TTL_SECONDS, +): Promise { + const result = await redis.writeRuntimeSessionRecordScript( + `${SESS_PREFIX}${record.runtime_session_id}`, + `${LOCK_PREFIX}${record.runtime_session_id}`, + lockToken, + JSON.stringify(record), + String(ttlSeconds), + ); + return result === 1; +} + +/** Monotonic generation for launch fencing: allocated while holding the lock, + * before RunMicrovm, so a stale worker's record can never outrank a newer + * launch. */ +export async function allocateRuntimeSessionGeneration(runtimeSessionId: string): Promise { + const key = `${GEN_PREFIX}${runtimeSessionId}`; + const generation = await redis.incr(key); + await redis.expire(key, RUNTIME_SESSION_RECORD_TTL_SECONDS); + return generation; +} + +/** Atomically reserves a monotonic checkpoint sequence above both the Redis + * counter and the highest object retained in durable storage. Combining the + * high-water reseed and increment in one Lua operation is required for + * fencing: a stale holder that resumes after another worker acquired the + * session lock can consume a distinct sequence, but can never reset the + * counter and overwrite the newer holder's immutable object key. */ +export async function allocateCheckpointSequence( + runtimeSessionId: string, + retainedMax = 0, +): Promise { + const key = `${CKPT_SEQ_PREFIX}${runtimeSessionId}`; + return redis.allocateCheckpointSequenceScript( + key, + String(retainedMax), + String(RUNTIME_SESSION_RECORD_TTL_SECONDS), + ); +} + +/** Fenced removal: deletes the record while the caller holds the mutex. + * Returns false when fenced. */ +export async function removeRuntimeSession(runtimeSessionId: string, lockToken: string): Promise { + const result = await redis.removeRuntimeSessionScript( + `${SESS_PREFIX}${runtimeSessionId}`, + `${LOCK_PREFIX}${runtimeSessionId}`, + lockToken, + ); + return result === 1; +} diff --git a/service/src/runtime-session/throttle.test.ts b/service/src/runtime-session/throttle.test.ts new file mode 100644 index 0000000..739b3c1 --- /dev/null +++ b/service/src/runtime-session/throttle.test.ts @@ -0,0 +1,103 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import { + MicrovmOpThrottledError, + acquireOpBudget, + poisonOpBucket, + resetRedisForTests, + setRedisForTests, +} from './throttle'; + +let mock: InstanceType; + +beforeEach(async () => { + /* ioredis-mock shares one keyspace across instances — flush per test. */ + mock = new RedisMock(); + await mock.flushall(); + setRedisForTests(mock); +}); + +afterEach(() => { + resetRedisForTests(); +}); + +/** Virtual clock: sleep() advances time instead of waiting. */ +function virtualClock(startMs = 1_000_000_000_000): { + now: () => number; + sleep: (ms: number) => Promise; + slept: number[]; +} { + let t = startMs; + const slept: number[] = []; + return { + now: () => t, + sleep: (ms: number) => { + slept.push(ms); + t += ms; + return Promise.resolve(); + }, + slept, + }; +} + +describe('acquireOpBudget', () => { + test('grants slots under the per-second limit without sleeping', async () => { + const clock = virtualClock(); + for (let i = 0; i < 4; i++) { + await acquireOpBudget('run', { limitPerSecond: 4, budgetMs: 5_000, now: clock.now, sleep: clock.sleep }); + } + expect(clock.slept).toHaveLength(0); + }); + + test('the (limit+1)th call in one second waits into the next second', async () => { + const clock = virtualClock(); + for (let i = 0; i < 5; i++) { + await acquireOpBudget('run', { limitPerSecond: 4, budgetMs: 5_000, now: clock.now, sleep: clock.sleep }); + } + expect(clock.slept.length).toBeGreaterThanOrEqual(1); + expect(clock.slept[0]).toBeGreaterThanOrEqual(1_000); + expect(clock.slept[0]).toBeLessThan(1_200); + }); + + test('throws MicrovmOpThrottledError when the budget cannot cover the wait', async () => { + const clock = virtualClock(); + for (let i = 0; i < 4; i++) { + await acquireOpBudget('run', { limitPerSecond: 4, budgetMs: 5_000, now: clock.now, sleep: clock.sleep }); + } + expect( + acquireOpBudget('run', { limitPerSecond: 4, budgetMs: 100, now: clock.now, sleep: clock.sleep }), + ).rejects.toThrow(MicrovmOpThrottledError); + }); + + test('ops have independent buckets', async () => { + const clock = virtualClock(); + await acquireOpBudget('suspend', { limitPerSecond: 1, budgetMs: 100, now: clock.now, sleep: clock.sleep }); + await acquireOpBudget('run', { limitPerSecond: 1, budgetMs: 100, now: clock.now, sleep: clock.sleep }); + expect(clock.slept).toHaveLength(0); + }); + + test('a poisoned bucket blocks until it clears, then grants', async () => { + const clock = virtualClock(); + await poisonOpBucket('run', 60_000); + expect( + acquireOpBudget('run', { limitPerSecond: 4, budgetMs: 200, now: clock.now, sleep: clock.sleep }), + ).rejects.toThrow(MicrovmOpThrottledError); + + await mock.del('rtsx:tps:poison:run'); + await acquireOpBudget('run', { limitPerSecond: 4, budgetMs: 200, now: clock.now, sleep: clock.sleep }); + }); + + test('cancellation interrupts a real throttle wait immediately', async () => { + await poisonOpBucket('run', 60_000); + const controller = new AbortController(); + const started = Date.now(); + setTimeout(() => controller.abort(new Error('job deadline')), 10); + + await expect(acquireOpBudget('run', { + limitPerSecond: 4, + budgetMs: 120_000, + signal: controller.signal, + })).rejects.toThrow('job deadline'); + expect(Date.now() - started).toBeLessThan(1_000); + }); +}); diff --git a/service/src/runtime-session/throttle.ts b/service/src/runtime-session/throttle.ts new file mode 100644 index 0000000..dc21cb4 --- /dev/null +++ b/service/src/runtime-session/throttle.ts @@ -0,0 +1,118 @@ +import type { Redis } from 'ioredis'; +import { connection } from '../queue'; + +/** + * Distributed per-second token buckets for Lambda MicroVM control-plane + * calls. All workers share the AWS account limits (RunMicrovm 5 TPS, + * ResumeMicrovm 5, SuspendMicrovm 2, CreateMicrovmAuthToken 50), so the + * budget lives in Redis: + * + * rtsx:tps:: INCR-ed per attempt (PEXPIRE 2s) + * rtsx:tps:poison: backoff flag set on SDK throttle errors + */ + +export type ThrottledOp = 'run' | 'resume' | 'suspend' | 'token'; + +const BUCKET_PREFIX = 'rtsx:tps:'; +const POISON_PREFIX = 'rtsx:tps:poison:'; +const BUCKET_TTL_MS = 2_000; +const DEFAULT_POISON_MS = 2_000; + +let redis: Redis = connection; + +export function setRedisForTests(client: Redis): void { + redis = client; +} + +export function resetRedisForTests(): void { + redis = connection; +} + +export class MicrovmOpThrottledError extends Error { + constructor(public readonly op: ThrottledOp, budgetMs: number) { + super(`Lambda MicroVM ${op} budget exhausted after ${budgetMs}ms of throttling`); + this.name = 'MicrovmOpThrottledError'; + } +} + +export interface OpBudgetOptions { + limitPerSecond: number; + /** Total time the caller is willing to wait for a slot. */ + budgetMs: number; + now?: () => number; + sleep?: (ms: number, signal?: AbortSignal) => Promise; + signal?: AbortSignal; +} + +function abortReason(signal: AbortSignal): Error { + return signal.reason instanceof Error + ? signal.reason + : new Error('Lambda MicroVM operation budget aborted'); +} + +const defaultSleep = (ms: number, signal?: AbortSignal): Promise => new Promise( + (resolve, reject) => { + if (signal?.aborted) { + reject(abortReason(signal)); + return; + } + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + const onAbort = (): void => { + clearTimeout(timer); + reject(abortReason(signal as AbortSignal)); + }; + signal?.addEventListener('abort', onAbort, { once: true }); + }, +); + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) throw abortReason(signal); +} + +/** + * Reserves one control-plane call slot for `op`, waiting across second + * boundaries until `budgetMs` is exhausted. Throws MicrovmOpThrottledError + * when no slot frees up in time. + */ +export async function acquireOpBudget(op: ThrottledOp, options: OpBudgetOptions): Promise { + const now = options.now ?? Date.now; + const sleep = options.sleep ?? defaultSleep; + const deadline = now() + options.budgetMs; + + for (;;) { + throwIfAborted(options.signal); + const poisoned = await redis.pttl(`${POISON_PREFIX}${op}`); + throwIfAborted(options.signal); + if (poisoned > 0) { + if (now() + poisoned > deadline) throw new MicrovmOpThrottledError(op, options.budgetMs); + await sleep(poisoned, options.signal); + continue; + } + + const nowMs = now(); + const second = Math.floor(nowMs / 1_000); + const key = `${BUCKET_PREFIX}${op}:${second}`; + const count = await redis.incr(key); + throwIfAborted(options.signal); + if (count === 1) { + await redis.pexpire(key, BUCKET_TTL_MS); + throwIfAborted(options.signal); + } + if (count <= options.limitPerSecond) return; + + const nextSecondMs = (second + 1) * 1_000 - nowMs; + const jitter = Math.floor(Math.random() * 100); + const waitMs = nextSecondMs + jitter; + if (nowMs + waitMs > deadline) throw new MicrovmOpThrottledError(op, options.budgetMs); + await sleep(waitMs, options.signal); + } +} + +/** Called when the SDK reports ThrottlingException/TooManyRequests: back the + * whole fleet off `op` briefly instead of hammering per-second buckets. */ +export async function poisonOpBucket(op: ThrottledOp, durationMs: number = DEFAULT_POISON_MS): Promise { + await redis.set(`${POISON_PREFIX}${op}`, '1', 'PX', durationMs); +} diff --git a/service/src/sandbox-backend/http.test.ts b/service/src/sandbox-backend/http.test.ts new file mode 100644 index 0000000..d337064 --- /dev/null +++ b/service/src/sandbox-backend/http.test.ts @@ -0,0 +1,149 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from 'bun:test'; +import axios from 'axios'; +import { env } from '../config'; +import { HttpSandboxBackend } from './http'; +import type { SandboxExecuteContext, SandboxTransportRequest } from './types'; +import type * as t from '../types'; + +type CapturedRequest = { + method: string; + path: string; + rawBody: string; + headers: Record; +}; + +let server: ReturnType; +let captured: CapturedRequest[] = []; +let nextResponse: { status: number; body: unknown; delayMs?: number } = { status: 200, body: {} }; + +const savedEndpoint = env.SANDBOX_ENDPOINT; + +beforeAll(() => { + server = Bun.serve({ + port: 0, + async fetch(req) { + captured.push({ + method: req.method, + path: new URL(req.url).pathname, + rawBody: await req.text(), + headers: Object.fromEntries(req.headers.entries()), + }); + if (nextResponse.delayMs) { + await new Promise((resolve) => setTimeout(resolve, nextResponse.delayMs)); + } + return new Response(JSON.stringify(nextResponse.body), { + status: nextResponse.status, + headers: { 'Content-Type': 'application/json' }, + }); + }, + }); + env.SANDBOX_ENDPOINT = `http://localhost:${server.port}/api/v2`; +}); + +afterAll(() => { + env.SANDBOX_ENDPOINT = savedEndpoint; + server.stop(true); +}); + +afterEach(() => { + captured = []; + nextResponse = { status: 200, body: {} }; +}); + +function payloadBody(): t.PayloadBody { + return { + language: 'python', + version: '3.14.4', + session_id: 'sess_exec_1', + output_session_id: 'sess_out_1', + files: [{ id: 'file_1', storage_session_id: 'sess_store_1', name: 'inputs/data.csv' }], + egress_grant: 'ceg1.iv.ct.tag', + execution_manifest: 'signed-manifest-token', + env_vars: { PTC_HISTORY_PATH: '/mnt/data/_ptc_history.json' }, + }; +} + +function request(): SandboxTransportRequest { + return { body: payloadBody(), headers: { 'Content-Type': 'application/json' } }; +} + +function context(overrides: Partial = {}): SandboxExecuteContext { + return { + executionId: 'exec_1', + language: 'python', + isSynthetic: false, + signal: new AbortController().signal, + runtimeSessionMode: 'stateless', + ...overrides, + }; +} + +describe('HttpSandboxBackend', () => { + test('POSTs the request body byte-identical to SANDBOX_ENDPOINT/execute', async () => { + const responseBody = { + session_id: 'sess_exec_1', + language: 'python', + version: '3.14.4', + files: [], + run: { + stdout: 'ok', stderr: '', code: 0, signal: null, output: 'ok', + memory: 1, message: null, status: null, cpu_time: 1, wall_time: 2, + }, + }; + nextResponse = { status: 200, body: responseBody }; + + const backend = new HttpSandboxBackend(); + const req = request(); + const result = await backend.execute(req, context()); + + expect(captured).toHaveLength(1); + expect(captured[0].method).toBe('POST'); + expect(captured[0].path).toBe('/api/v2/execute'); + expect(captured[0].rawBody).toBe(JSON.stringify(req.body)); + expect(captured[0].headers['content-type']).toBe('application/json'); + expect(result).toEqual(responseBody); + }); + + test('does not mutate the signed request body', async () => { + const req = request(); + const before = JSON.stringify(req.body); + await new HttpSandboxBackend().execute(req, context()); + expect(JSON.stringify(req.body)).toBe(before); + }); + + test('throws "Error from sandbox" on 2xx statuses other than 200', async () => { + nextResponse = { status: 201, body: { session_id: 'x' } }; + expect(new HttpSandboxBackend().execute(request(), context())) + .rejects.toThrow('Error from sandbox'); + }); + + test('rethrows axios errors untouched on non-2xx statuses', async () => { + nextResponse = { status: 500, body: { message: 'sandbox exploded' } }; + try { + await new HttpSandboxBackend().execute(request(), context()); + throw new Error('expected rejection'); + } catch (error) { + expect(axios.isAxiosError(error)).toBe(true); + if (axios.isAxiosError(error)) { + expect(error.response?.status).toBe(500); + expect(error.response?.data).toEqual({ message: 'sandbox exploded' }); + } + } + }); + + test('propagates the worker abort signal as an axios cancellation', async () => { + nextResponse = { status: 200, body: { session_id: 'x' }, delayMs: 5_000 }; + const controller = new AbortController(); + const pending = new HttpSandboxBackend().execute(request(), context({ signal: controller.signal })); + setTimeout(() => controller.abort(), 20); + try { + await pending; + throw new Error('expected rejection'); + } catch (error) { + expect(axios.isAxiosError(error)).toBe(true); + if (axios.isAxiosError(error)) { + expect(error.code === 'ERR_CANCELED' || error.name === 'AbortError').toBe(true); + } + } + }); +}); diff --git a/service/src/sandbox-backend/http.ts b/service/src/sandbox-backend/http.ts new file mode 100644 index 0000000..ce5f5c0 --- /dev/null +++ b/service/src/sandbox-backend/http.ts @@ -0,0 +1,34 @@ +import axios from 'axios'; +import type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +import { injectTraceHeaders, withSpan } from '../telemetry'; +import { Jobs } from '../enum'; +import { env } from '../config'; + +/** Current behavior: POST the signed request to SANDBOX_ENDPOINT. + * Axios errors are rethrown untouched so the worker's existing + * abort/timeout/sandbox-error mapping stays byte-identical. */ +export class HttpSandboxBackend implements SandboxBackend { + readonly name = 'http' as const; + + async execute(req: SandboxTransportRequest, ctx: SandboxExecuteContext): Promise { + const response = await withSpan('codeapi.sandbox.execute', { + 'http.request.method': 'POST', + 'url.path': `/${Jobs.execute}`, + 'codeapi.language': ctx.language, + 'codeapi.sandbox.backend': this.name, + }, () => axios.post( + `${env.SANDBOX_ENDPOINT}/${Jobs.execute}`, + req.body, + { + headers: injectTraceHeaders(req.headers), + signal: ctx.signal, + } + ), 'CLIENT'); + + if (response.status !== 200) { + throw new Error('Error from sandbox'); + } + + return response.data; + } +} diff --git a/service/src/sandbox-backend/index.test.ts b/service/src/sandbox-backend/index.test.ts new file mode 100644 index 0000000..e9f2ac7 --- /dev/null +++ b/service/src/sandbox-backend/index.test.ts @@ -0,0 +1,71 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import * as path from 'path'; +import { env } from '../config'; +import { HttpSandboxBackend } from './http'; +import { getSandboxBackend, setSandboxBackendForTests } from './index'; +import type { SandboxBackend } from './types'; + +const savedBackend = env.SANDBOX_BACKEND; + +afterEach(() => { + env.SANDBOX_BACKEND = savedBackend; + setSandboxBackendForTests(undefined); +}); + +describe('getSandboxBackend', () => { + test('defaults to the http backend and memoizes it', () => { + const backend = getSandboxBackend(); + expect(backend).toBeInstanceOf(HttpSandboxBackend); + expect(backend.name).toBe('http'); + expect(getSandboxBackend()).toBe(backend); + }); + + test('selects the lazy lambda-microvm backend when configured', () => { + env.SANDBOX_BACKEND = 'lambda-microvm'; + const backend = getSandboxBackend(); + expect(backend.name).toBe('lambda-microvm'); + }); + + test('does not load Lambda-only modules for the HTTP backend', async () => { + const serviceRoot = path.resolve(import.meta.dir, '../..'); + const probe = Bun.spawn([ + process.execPath, + '-e', + ` + process.env.CODEAPI_SANDBOX_BACKEND = 'http'; + const { getSandboxBackend } = await import('./src/sandbox-backend/index.ts'); + if (getSandboxBackend().name !== 'http') process.exit(2); + const loaded = Object.keys(require.cache).filter((id) => + id.includes('/sandbox-backend/lambda-microvm.') + || id.includes('/runtime-session/checkpoint-store.') + || id.includes('/@aws-sdk/client-s3/') + || id.includes('/@aws-sdk/client-lambda-microvms/') + ); + console.log(JSON.stringify(loaded)); + `, + ], { + cwd: serviceRoot, + env: { ...process.env, CODEAPI_SANDBOX_BACKEND: 'http' }, + stdout: 'pipe', + stderr: 'pipe', + }); + const [exitCode, stdout, stderr] = await Promise.all([ + probe.exited, + new Response(probe.stdout).text(), + new Response(probe.stderr).text(), + ]); + expect(exitCode, stderr).toBe(0); + expect(JSON.parse(stdout.trim().split('\n').at(-1) ?? 'null')).toEqual([]); + }); + + test('test seam replaces the active backend', () => { + const fake: SandboxBackend = { + name: 'http', + execute: () => Promise.reject(new Error('unused')), + }; + setSandboxBackendForTests(fake); + expect(getSandboxBackend()).toBe(fake); + setSandboxBackendForTests(undefined); + expect(getSandboxBackend()).toBeInstanceOf(HttpSandboxBackend); + }); +}); diff --git a/service/src/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts new file mode 100644 index 0000000..3ebf796 --- /dev/null +++ b/service/src/sandbox-backend/index.ts @@ -0,0 +1,91 @@ +import type { + SandboxBackend, + SandboxExecuteContext, + SandboxRawResponse, + SandboxTransportRequest, +} from './types'; +import type { LambdaMicrovmClient } from '../runtime-session/lambda-client'; +import { HttpSandboxBackend } from './http'; +import { env } from '../config'; + +export type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +export { SandboxBackendError } from './types'; +export { HttpSandboxBackend } from './http'; + +let backend: SandboxBackend | undefined; + +class LazyLambdaMicrovmSandboxBackend implements SandboxBackend { + readonly name = 'lambda-microvm' as const; + private backendPromise: Promise | undefined; + + private load(): Promise { + this.backendPromise ??= (async (): Promise => { + const { LambdaMicrovmSandboxBackend } = await import('./lambda-microvm'); + const checkpointStore = env.SESSION_CHECKPOINTS && env.RUNTIME_SESSION_MODE !== 'stateless' + ? new (await import('../runtime-session/checkpoint-store')).MinioCheckpointStore() + : undefined; + + return new LambdaMicrovmSandboxBackend({ + /* The AWS client stays behind the same backend gate. */ + clientFactory: async (): Promise => { + const { AwsLambdaMicrovmClient } = await import('../runtime-session/lambda-client-aws'); + return new AwsLambdaMicrovmClient({ region: env.LAMBDA_MICROVM_REGION }); + }, + config: { + imageArn: env.LAMBDA_MICROVM_IMAGE_ARN, + imageVersion: env.LAMBDA_MICROVM_IMAGE_VERSION, + executionRoleArn: env.LAMBDA_MICROVM_EXECUTION_ROLE_ARN, + logGroup: env.LAMBDA_MICROVM_LOG_GROUP, + ingressConnectorArns: env.LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS, + egressConnectorArns: env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS, + port: env.LAMBDA_MICROVM_PORT, + maxDurationSeconds: env.LAMBDA_MICROVM_MAX_DURATION_SECONDS, + authTokenTtlSeconds: env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS, + launchTimeoutMs: env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, + healthTimeoutMs: env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS, + launchTps: env.LAMBDA_MICROVM_LAUNCH_TPS, + tokenTps: env.LAMBDA_MICROVM_TOKEN_TPS, + jobTimeoutMs: env.JOB_TIMEOUT, + idleSeconds: env.LAMBDA_MICROVM_IDLE_SECONDS, + suspendedSeconds: env.LAMBDA_MICROVM_SUSPEND_SECONDS, + lockWaitMs: env.RUNTIME_SESSION_LOCK_WAIT_MS, + checkpointsEnabled: env.SESSION_CHECKPOINTS, + checkpoint: { + port: env.LAMBDA_MICROVM_PORT, + authTokenTtlSeconds: env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS, + maxBytes: env.CHECKPOINT_MAX_BYTES, + timeoutMs: env.CHECKPOINT_TIMEOUT_MS, + }, + }, + checkpointStore, + }); + })(); + return this.backendPromise; + } + + async execute( + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + ): Promise { + return (await this.load()).execute(req, ctx); + } +} + +function createBackend(): SandboxBackend { + if (env.SANDBOX_BACKEND === 'lambda-microvm') { + /* Loading the concrete backend also loads its session registry and + * checkpoint code. Defer the whole graph so the default HTTP worker does + * not require AWS SDK modules or initialize Lambda-only dependencies. */ + return new LazyLambdaMicrovmSandboxBackend(); + } + return new HttpSandboxBackend(); +} + +export function getSandboxBackend(): SandboxBackend { + backend ??= createBackend(); + return backend; +} + +export function setSandboxBackendForTests(next: SandboxBackend | undefined): void { + backend = next; +} diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts new file mode 100644 index 0000000..042f670 --- /dev/null +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -0,0 +1,1451 @@ +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import * as zlib from 'zlib'; +import * as fsp from 'fs/promises'; +import axios from 'axios'; +import { env } from '../config'; +import { FakeLambdaMicrovmClient } from '../runtime-session/lambda-client-fake'; +import { LambdaMicrovmApiError } from '../runtime-session/lambda-client'; +import { + resetRedisForTests as resetThrottleRedis, + setRedisForTests as setThrottleRedis, +} from '../runtime-session/throttle'; +import { + acquireRuntimeSessionLock, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + resetRedisForTests as resetRegistryRedis, + setRedisForTests as setRegistryRedis, + writeRuntimeSessionRecord, +} from '../runtime-session/registry'; +import { MemoryCheckpointStore, checkpointObjectKey, checkpointPrefixFor } from '../runtime-session/checkpoint-store'; +import { + LambdaMicrovmSandboxBackend, + normalizeMicrovmEndpoint, + runtimeSessionLaunchFingerprint, + type LambdaMicrovmBackendConfig, +} from './lambda-microvm'; +import { SandboxBackendError } from './types'; +import type { SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +import type * as t from '../types'; + +type CapturedRequest = { path: string; rawBody: string; headers: Record }; + +let server: ReturnType; +let captured: CapturedRequest[] = []; +let healthStatus = 200; +let executeDelayMs = 0; +let executeStatus = 200; +let executeResponseBody: unknown; +let sessionFilesStatus = 200; +let lastSessionFilesBody: Buffer | null = null; +let stealSessionLockOnExecute = false; +let fileReadOnly = false; +let fileObjectStatus = 200; +let probeMissingOverride: Array<{ cache_key: string }> | undefined; +let evictAfterPush = false; +let recordDuringRestore: Awaited> | undefined; +let onExecute: (() => void | Promise) | undefined; +/** Models the runner's input cache so probe/push behave like the real VM. */ +let lastProbedRefs: Array<{ cache_key: string }> = []; +const vmInputCache = new Set(); +const fileObjectBytes = 'csv,bytes\n1,2\n'; +let mock: InstanceType; +const checkpointBlob = 'FAKE_TAR_GZ_BYTES'; + +const EXECUTE_RESPONSE = { + session_id: 'sess_exec_1', + language: 'python', + version: '3.14.4', + files: [], + run: { + stdout: 'ok', stderr: '', code: 0, signal: null, output: 'ok', + memory: 1, message: null, status: null, cpu_time: 1, wall_time: 2, + }, +}; + +beforeAll(() => { + server = Bun.serve({ + port: 0, + async fetch(req) { + const path = new URL(req.url).pathname; + const raw = Buffer.from(await req.arrayBuffer()); + captured.push({ + path, + rawBody: raw.toString(), + headers: Object.fromEntries(req.headers.entries()), + }); + /* The same server doubles as the internal file server the control plane + * fetches input refs from when building a session files delivery. */ + if (path.startsWith('/sessions/')) { + return new Response(fileObjectBytes, { + status: fileObjectStatus, + headers: fileReadOnly ? { 'X-Read-Only': 'true' } : {}, + }); + } + if (path === '/api/v2/session/inputs/probe') { + const refs = (JSON.parse(raw.toString()) as { + refs: Array<{ cache_key: string }>; + }).refs; + lastProbedRefs = refs; + const missing = probeMissingOverride + ?? refs.filter((r) => !vmInputCache.has(r.cache_key)); + return new Response(JSON.stringify({ missing }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + if (path === '/api/v2/session/inputs') { + lastSessionFilesBody = raw; + if (sessionFilesStatus === 200) { + /* Model the runner cache: everything just pushed is now held. */ + for (const ref of lastProbedRefs) vmInputCache.add(ref.cache_key); + if (evictAfterPush) vmInputCache.clear(); + } + return new Response(JSON.stringify({ stored: lastProbedRefs.length }), { + status: sessionFilesStatus, + headers: { 'Content-Type': 'application/json' }, + }); + } + if (path === '/api/v2/health') { + return new Response(JSON.stringify({ status: 'ok' }), { + status: healthStatus, + headers: { 'Content-Type': 'application/json' }, + }); + } + if (path === '/api/v2/execute') { + await onExecute?.(); + if (executeDelayMs > 0) { + await new Promise((resolve) => setTimeout(resolve, executeDelayMs)); + } + if (stealSessionLockOnExecute) { + await mock.set('rtsx:lock:rt_session_1', 'stolen'); + } + return new Response(JSON.stringify(executeResponseBody), { + status: executeStatus, + headers: { 'Content-Type': 'application/json' }, + }); + } + if (path === '/api/v2/session/checkpoint') { + return new Response(checkpointBlob, { status: 200, headers: { 'Content-Type': 'application/x-gtar' } }); + } + if (path === '/api/v2/session/restore') { + recordDuringRestore = await readRuntimeSessionRecord('rt_ckpt_1'); + return new Response(JSON.stringify({ status: 'restored' }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response('not found', { status: 404 }); + }, + }); + env.FILE_SERVER_URL = `http://localhost:${server.port}`; +}); + +afterAll(() => { + server.stop(true); +}); + +beforeEach(async () => { + /* ioredis-mock shares one keyspace across instances — flush per test. */ + mock = new RedisMock(); + await mock.flushall(); + setThrottleRedis(mock); + setRegistryRedis(mock); + captured = []; + healthStatus = 200; + executeDelayMs = 0; + executeStatus = 200; + executeResponseBody = EXECUTE_RESPONSE; + sessionFilesStatus = 200; + lastSessionFilesBody = null; + stealSessionLockOnExecute = false; + fileReadOnly = false; + fileObjectStatus = 200; + probeMissingOverride = undefined; + evictAfterPush = false; + recordDuringRestore = undefined; + onExecute = undefined; + lastProbedRefs = []; + vmInputCache.clear(); +}); + +afterEach(() => { + resetThrottleRedis(); + resetRegistryRedis(); +}); + +function config(overrides: Partial = {}): LambdaMicrovmBackendConfig { + return { + imageArn: 'arn:aws:lambda:us-east-2:1:microvm-image:codeapi', + imageVersion: '3', + port: 8080, + maxDurationSeconds: 28_800, + authTokenTtlSeconds: 300, + launchTimeoutMs: 2_000, + healthTimeoutMs: 1_000, + launchTps: 50, + tokenTps: 50, + jobTimeoutMs: 300_000, + idleSeconds: 300, + suspendedSeconds: 1_800, + lockWaitMs: 500, + checkpointsEnabled: false, + checkpoint: { port: 8080, authTokenTtlSeconds: 300, maxBytes: 512 * 1024 * 1024, timeoutMs: 30_000 }, + ...overrides, + }; +} + +function makeBackend( + fake: FakeLambdaMicrovmClient, + cfg?: Partial, + checkpointStore?: MemoryCheckpointStore, +): LambdaMicrovmSandboxBackend { + return new LambdaMicrovmSandboxBackend({ + clientFactory: () => Promise.resolve(fake), + config: config(cfg), + pollIntervalMs: 5, + checkpointStore, + }); +} + +function fakeClient(): FakeLambdaMicrovmClient { + return new FakeLambdaMicrovmClient({ endpointProvider: () => `http://localhost:${server.port}` }); +} + +function payloadBody(): t.PayloadBody { + return { + language: 'python', + version: '3.14.4', + session_id: 'sess_exec_1', + files: [{ id: 'file_1', storage_session_id: 'sess_store_1', name: 'inputs/data.csv' }], + egress_grant: 'ceg1.iv.ct.tag', + execution_manifest: 'signed-manifest-token', + }; +} + +function request(): SandboxTransportRequest { + return { body: payloadBody(), headers: { 'Content-Type': 'application/json' } }; +} + +function context(overrides: Partial = {}): SandboxExecuteContext { + return { + executionId: 'exec_42', + language: 'python', + isSynthetic: false, + signal: new AbortController().signal, + runtimeSessionMode: 'stateless', + ...overrides, + }; +} + +describe('normalizeMicrovmEndpoint', () => { + test('prefixes https for bare hosts and keeps explicit schemes', () => { + expect(normalizeMicrovmEndpoint('abc.lambda-microvm.on.aws')).toBe('https://abc.lambda-microvm.on.aws'); + expect(normalizeMicrovmEndpoint('abc.on.aws/')).toBe('https://abc.on.aws'); + expect(normalizeMicrovmEndpoint('http://localhost:1234')).toBe('http://localhost:1234'); + expect(normalizeMicrovmEndpoint('https://x.on.aws///')).toBe('https://x.on.aws'); + }); +}); + +describe('LambdaMicrovmSandboxBackend stateless execution', () => { + test('run -> health -> execute -> terminate happy path', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + const req = request(); + + const result = await backend.execute(req, context()); + + expect(result).toEqual(EXECUTE_RESPONSE); + + const runCalls = fake.callsFor('runMicrovm'); + expect(runCalls).toHaveLength(1); + const runArgs = runCalls[0].args as { imageIdentifier: string; clientToken?: string; maximumDurationSeconds: number }; + expect(runArgs.imageIdentifier).toBe('arn:aws:lambda:us-east-2:1:microvm-image:codeapi'); + expect(runArgs.clientToken).toBe('exec-exec_42'); + expect(runArgs.maximumDurationSeconds).toBe(Math.ceil(300_000 / 1_000) + 120); + + const executeReq = captured.find((c) => c.path === '/api/v2/execute'); + expect(executeReq).toBeDefined(); + expect(executeReq?.rawBody).toBe(JSON.stringify(req.body)); + const vm = [...fake.vms.values()][0]; + /* Input delivery mints its own tokens before the execute, so assert the + * execute carried one of THIS VM's tokens rather than a fixed index. */ + expect(vm.mintedTokens).toContain(executeReq?.headers['x-aws-proxy-auth'] as string); + + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(vm.state).toBe('TERMINATED'); + }); + + test('health check runs before execute', async () => { + const fake = fakeClient(); + await makeBackend(fake).execute(request(), context()); + const paths = captured.map((c) => c.path); + expect(paths.indexOf('/api/v2/health')).toBeGreaterThanOrEqual(0); + expect(paths.indexOf('/api/v2/health')).toBeLessThan(paths.indexOf('/api/v2/execute')); + }); + + test('health check runs before stateless input delivery', async () => { + const fake = fakeClient(); + await makeBackend(fake).execute(request(), context()); + const paths = captured.map((c) => c.path); + expect(paths.indexOf('/api/v2/health')).toBeLessThan( + paths.indexOf('/api/v2/session/inputs/probe'), + ); + }); + + test('terminates the VM even when the execute is aborted mid-flight', async () => { + const fake = fakeClient(); + executeDelayMs = 5_000; + const controller = new AbortController(); + const pending = makeBackend(fake).execute(request(), context({ signal: controller.signal })); + setTimeout(() => controller.abort(), 50); + + try { + await pending; + throw new Error('expected rejection'); + } catch (error) { + expect(axios.isAxiosError(error)).toBe(true); + } + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + test('launch poll timeout surfaces MICROVM_LAUNCH_FAILED and terminates the stuck VM', async () => { + const fake = fakeClient(); + fake.delayNextLaunch(10_000); + const backend = makeBackend(fake, { launchTimeoutMs: 60 }); + + try { + await backend.execute(request(), context()); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('MICROVM_LAUNCH_FAILED'); + } + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + test('control-plane throttle surfaces MICROVM_LAUNCH_THROTTLED and poisons the run bucket', async () => { + const fake = fakeClient(); + fake.failNext('runMicrovm', new LambdaMicrovmApiError('throttled', 'RunMicrovm', 'rate exceeded')); + + try { + await makeBackend(fake).execute(request(), context()); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('MICROVM_LAUNCH_THROTTLED'); + } + expect(await mock.exists('rtsx:tps:poison:run')).toBe(1); + }); + + test('failed health check surfaces MICROVM_UNHEALTHY and terminates', async () => { + const fake = fakeClient(); + healthStatus = 500; + + try { + await makeBackend(fake).execute(request(), context()); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('MICROVM_UNHEALTHY'); + } + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + test('refreshes an expiring proxy token while waiting for runner readiness', async () => { + const fake = fakeClient(); + const mint = fake.createMicrovmAuthToken.bind(fake); + let readinessMints = 0; + healthStatus = 500; + fake.createMicrovmAuthToken = async (args) => { + readinessMints += 1; + const token = await mint(args); + if (readinessMints === 1) { + return { ...token, expiresAtMs: Date.now() + 1 }; + } + /* A second readiness mint is what makes the simulated runner ready. + * Without rollover the first 500 repeats until launch timeout. */ + if (readinessMints === 2) healthStatus = 200; + return token; + }; + + await expect( + makeBackend(fake, { launchTimeoutMs: 250, healthTimeoutMs: 20 }) + .execute(request(), context()), + ).resolves.toEqual(EXECUTE_RESPONSE); + expect(readinessMints).toBeGreaterThanOrEqual(2); + }); + + test('does not mutate the signed request body', async () => { + const fake = fakeClient(); + const req = request(); + const before = JSON.stringify(req.body); + await makeBackend(fake).execute(req, context()); + expect(JSON.stringify(req.body)).toBe(before); + }); + + test('a MicroVM that dies during boot is retried once with a fresh clientToken', async () => { + const fake = fakeClient(); + fake.terminateNextLaunch(); + + const result = await makeBackend(fake).execute(request(), context()); + + expect(result).toEqual(EXECUTE_RESPONSE); + const runCalls = fake.callsFor('runMicrovm'); + expect(runCalls).toHaveLength(2); + const tokens = runCalls.map((call) => (call.args as { clientToken?: string }).clientToken); + expect(tokens[0]).toBe('exec-exec_42'); + expect(tokens[1]).toBe('exec-exec_42-r1'); + }); + + test('a second boot-time death fails the request — single retry only', async () => { + const fake = fakeClient(); + fake.terminateNextLaunch(); + fake.terminateNextLaunch(); + + await expect(makeBackend(fake).execute(request(), context())).rejects.toMatchObject({ + code: 'MICROVM_LAUNCH_FAILED', + }); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + }); +}); + +describe('LambdaMicrovmSandboxBackend session execution', () => { + test('retries transport failures while a suspended input endpoint resumes', async () => { + const reservation = Bun.serve({ + hostname: '127.0.0.1', + port: 0, + fetch: () => new Response('reserved'), + }); + const port = reservation.port; + reservation.stop(true); + + let lateServer: ReturnType | undefined; + const starter = setTimeout(() => { + lateServer = Bun.serve({ + hostname: '127.0.0.1', + port, + fetch: () => new Response(JSON.stringify({ missing: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }), + }); + }, 25); + + const backend = makeBackend(fakeClient()) as unknown as { + probeInputsWithRetry( + args: { + mintToken: () => Promise<{ + headerName: string; + token: string; + expiresAtMs: number; + }>; + endpointBase: string; + signal?: AbortSignal; + }, + refs: Array<{ cache_key: string }>, + cfg: { + port: number; + authTokenTtlSeconds: number; + maxBytes: number; + timeoutMs: number; + }, + ): Promise>; + }; + + try { + await expect(backend.probeInputsWithRetry( + { + mintToken: async () => ({ + headerName: 'X-aws-proxy-auth', + token: 'test-token', + expiresAtMs: Date.now() + 60_000, + }), + endpointBase: `http://127.0.0.1:${port}`, + }, + [{ cache_key: 'a'.repeat(64) }], + { + port: 8080, + authTokenTtlSeconds: 300, + maxBytes: 1024, + timeoutMs: 500, + }, + )).resolves.toEqual([]); + } finally { + clearTimeout(starter); + lateServer?.stop(true); + } + }); + + function sessionContext(overrides: Partial = {}): SandboxExecuteContext { + return context({ + runtimeSessionId: 'rt_session_1', + runtimeSessionMode: 'affinity', + tenantId: 'tenant-a', + canonicalUserId: 'user-1', + ...overrides, + }); + } + + test('launches a hookless session VM (idlePolicy, no runHookPayload) and stamps the workspace header on execute', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + + const result = await backend.execute(request(), sessionContext()); + expect(result).toEqual(EXECUTE_RESPONSE); + + const runArgs = fake.callsFor('runMicrovm')[0].args as { + runHookPayload?: string; + idlePolicy?: { autoResume: boolean; maxIdleSeconds: number }; + clientToken?: string; + maximumDurationSeconds: number; + }; + /* Session mode is delivered per-request via the header, never a /run hook + * (image builds stay hookless), so RunMicrovm carries no runHookPayload. */ + expect(runArgs.runHookPayload).toBeUndefined(); + expect(runArgs.idlePolicy?.autoResume).toBe(true); + expect(runArgs.clientToken).toBe('sess-rt_session_1-1'); + expect(runArgs.maximumDurationSeconds).toBe(28_800); + + const executeReq = captured.find((c) => c.path === '/api/v2/execute'); + expect(executeReq?.headers['x-runtime-session-id']).toBe('rt_session_1'); + + const record = await readRuntimeSessionRecord('rt_session_1'); + expect(record?.state).toBe('RUNNING'); + expect(record?.microvm_id).toBe([...fake.vms.keys()][0]); + expect(record?.generation).toBe(1); + }); + + test('replays a recorded launch intent after the RunMicrovm response is lost', async () => { + const fake = fakeClient(); + const cfg = config(); + const launchedAt = Date.now(); + const token = 'sess-rt_session_1-7'; + const accepted = await fake.runMicrovm({ + imageIdentifier: cfg.imageArn, + imageVersion: cfg.imageVersion, + executionRoleArn: cfg.executionRoleArn, + logGroup: cfg.logGroup, + ingressConnectorArns: cfg.ingressConnectorArns, + egressConnectorArns: cfg.egressConnectorArns, + maximumDurationSeconds: cfg.maxDurationSeconds, + idlePolicy: { + maxIdleSeconds: cfg.idleSeconds, + suspendedSeconds: cfg.suspendedSeconds, + autoResume: true, + }, + clientToken: token, + }); + /* Model a worker that persisted its intent and whose RunMicrovm reached AWS, + * but died before it could record the returned MicroVM id. */ + await mock.set('rtsx:gen:rt_session_1', '7'); + const lock = await acquireRuntimeSessionLock('rt_session_1', 60_000); + expect(lock).not.toBeNull(); + await writeRuntimeSessionRecord({ + runtime_session_id: 'rt_session_1', + tenant_id: 'tenant-a', + canonical_user_id: 'user-1', + port: cfg.port, + image_arn: cfg.imageArn, + image_version: cfg.imageVersion, + launch_fingerprint: runtimeSessionLaunchFingerprint(cfg), + state: 'PENDING', + generation: 7, + launched_at: launchedAt, + last_seen_at: launchedAt, + hard_deadline_at: launchedAt + cfg.maxDurationSeconds * 1_000 - 60_000, + }, lock as string); + await releaseRuntimeSessionLock('rt_session_1', lock as string); + + const result = await makeBackend(fake).execute(request(), sessionContext()); + expect(result).toEqual(EXECUTE_RESPONSE); + const runCalls = fake.callsFor('runMicrovm'); + expect(runCalls).toHaveLength(2); + expect(runCalls.map(call => (call.args as { clientToken?: string }).clientToken)) + .toEqual([token, token]); + expect(fake.vms.size).toBe(1); + expect((await readRuntimeSessionRecord('rt_session_1'))?.microvm_id) + .toBe(accepted.microvmId); + }); + + test('reuses the warm VM on the second execution (no second RunMicrovm)', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + + await backend.execute(request(), sessionContext()); + await backend.execute(request(), sessionContext()); + + expect(fake.callsFor('runMicrovm')).toHaveLength(1); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + const executes = captured.filter((c) => c.path === '/api/v2/execute'); + expect(executes).toHaveLength(2); + }); + + test('a reused VM skips the preflight health check so a slow auto-resume can proceed', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + captured = []; + await backend.execute(request(), sessionContext()); + /* The warm/suspended VM auto-resumes on the execute itself under the full + * job budget; a 5s health probe would misclassify a slow resume as + * unhealthy and tear the VM down. */ + expect(captured.filter((c) => c.path === '/api/v2/health')).toHaveLength(0); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(1); + }); + + test('a runner non-2xx keeps the warm VM (does not tear down the session)', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + executeStatus = 500; + /* The runner responded (500) — the VM is alive, only the request failed — + * so the session must NOT be terminated (regression: the error-classifier + * previously tore down any error that wasn't literally "Error from sandbox"). */ + await expect(backend.execute(request(), sessionContext())).rejects.toThrow(); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + const record = await readRuntimeSessionRecord('rt_session_1'); + expect(record?.state).toBe('RUNNING'); + }); + + test('a partial-prime signal recycles the dirty session workspace', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + executeStatus = 409; + executeResponseBody = { + error: 'session_workspace_dirty', + message: 'Session workspace must be restored', + }; + + try { + await backend.execute(request(), sessionContext()); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('MICROVM_UNHEALTHY'); + } + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + const recycled = await readRuntimeSessionRecord('rt_session_1'); + expect(recycled?.state).toBe('TERMINATED'); + expect(recycled?.microvm_id).toBeUndefined(); + + executeStatus = 200; + executeResponseBody = EXECUTE_RESPONSE; + await expect(backend.execute(request(), sessionContext())).resolves.toEqual(EXECUTE_RESPONSE); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + }); + + + + + + + + + test('probes the VM and pushes only what it is missing', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + + const paths = captured.map((c) => c.path); + expect(paths.indexOf('/api/v2/session/inputs/probe')).toBeGreaterThanOrEqual(0); + expect(paths.indexOf('/api/v2/session/inputs')).toBeGreaterThan( + paths.indexOf('/api/v2/session/inputs/probe'), + ); + expect(paths.indexOf('/api/v2/session/inputs')).toBeLessThan(paths.indexOf('/api/v2/execute')); + const push = captured.find((request) => request.path === '/api/v2/session/inputs'); + expect(Number(push?.headers['x-codeapi-input-expanded-bytes'])).toBe( + Buffer.byteLength(fileObjectBytes) + + Buffer.byteLength(JSON.stringify({ readOnly: false })), + ); + /* Objects are pushed under runner-computed digests, and carry only + * object-level metadata: no caller-supplied path appears anywhere in the + * batch, so nothing in delivery can act on one. */ + const untarred = zlib.gunzipSync(lastSessionFilesBody!).toString('latin1'); + expect(untarred).toMatch(/[0-9a-f]{64}/); + expect(untarred).not.toContain('inputs/data.csv'); + }); + + test('a second execution pushes nothing when the VM already holds the object', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + captured = []; + await backend.execute(request(), sessionContext()); + + /* Dedupe is the VM's answer, not control-plane bookkeeping — so it stays + * correct across record loss, and a re-push could never revert an edit + * anyway (this path does not write the workspace). */ + expect(captured.filter((c) => c.path === '/api/v2/session/inputs/probe')).toHaveLength(1); + expect(captured.filter((c) => c.path === '/api/v2/session/inputs')).toHaveLength(0); + }); + + test('a stateless one-shot receives its by-ref inputs too', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), context()); + + /* The cache is keyed by object, not session, so the same mechanism serves + * stateless execution — which previously ran with its inputs missing. */ + const paths = captured.map((c) => c.path); + expect(paths.indexOf('/api/v2/session/inputs')).toBeGreaterThanOrEqual(0); + expect(paths.indexOf('/api/v2/session/inputs')).toBeLessThan(paths.indexOf('/api/v2/execute')); + }); + + test('a live runner rejecting an input push keeps the warm session VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + sessionFilesStatus = 500; + + await expect(backend.execute(request(), sessionContext())).rejects.toThrow( + 'Session input push failed', + ); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + expect((await readRuntimeSessionRecord('rt_session_1'))?.state).toBe('RUNNING'); + }); + + test('a proxy 502 during input push recycles the unreachable session VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + sessionFilesStatus = 502; + + await expect(backend.execute(request(), sessionContext())).rejects.toThrow( + 'Session input push failed', + ); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); + }); + + test('a probe token throttle preserves its code and keeps the warm session VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + vmInputCache.clear(); + captured = []; + fake.failNext( + 'createMicrovmAuthToken', + new LambdaMicrovmApiError('throttled', 'CreateMicrovmAuthToken', 'rate exceeded'), + ); + + await expect(backend.execute(request(), sessionContext())).rejects.toMatchObject({ + code: 'MICROVM_LAUNCH_THROTTLED', + }); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + expect((await readRuntimeSessionRecord('rt_session_1'))?.state).toBe('RUNNING'); + }); + + test('a push token throttle preserves its code and keeps the warm session VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + vmInputCache.clear(); + captured = []; + + const mint = fake.createMicrovmAuthToken.bind(fake); + let deliveryMints = 0; + fake.createMicrovmAuthToken = async (args) => { + deliveryMints += 1; + if (deliveryMints === 2) { + throw new LambdaMicrovmApiError( + 'throttled', + 'CreateMicrovmAuthToken', + 'rate exceeded', + ); + } + return mint(args); + }; + + await expect(backend.execute(request(), sessionContext())).rejects.toMatchObject({ + code: 'MICROVM_LAUNCH_THROTTLED', + }); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + expect((await readRuntimeSessionRecord('rt_session_1'))?.state).toBe('RUNNING'); + }); + + test('a source-object failure never recycles a healthy warm session VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + vmInputCache.clear(); + fileObjectStatus = 404; + captured = []; + + await expect(backend.execute(request(), sessionContext())).rejects.toMatchObject({ + code: 'SESSION_INPUT_UNAVAILABLE', + }); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + expect((await readRuntimeSessionRecord('rt_session_1'))?.state).toBe('RUNNING'); + }); + + test('an oversized input returns a typed limit error and keeps the warm session VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake, { + checkpoint: { ...config().checkpoint, maxBytes: 1 }, + }); + const warmup = request(); + warmup.body = { ...warmup.body, files: [] }; + await backend.execute(warmup, sessionContext()); + captured = []; + + await expect(backend.execute(request(), sessionContext())).rejects.toMatchObject({ + code: 'SESSION_INPUT_TOO_LARGE', + }); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + expect((await readRuntimeSessionRecord('rt_session_1'))?.state).toBe('RUNNING'); + }); + + test('rejects a probe response containing an unrequested cache key', async () => { + const fake = fakeClient(); + probeMissingOverride = [{ cache_key: 'f'.repeat(64) }]; + + await expect(makeBackend(fake).execute(request(), sessionContext())).rejects.toThrow( + 'Session input delivery failed', + ); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(0); + }); + + test('re-probes the full working set after push and never executes after eviction', async () => { + const fake = fakeClient(); + evictAfterPush = true; + + await expect(makeBackend(fake).execute(request(), sessionContext())).rejects.toThrow( + 'working set', + ); + const paths = captured.map((c) => c.path); + const pushIndex = paths.indexOf('/api/v2/session/inputs'); + expect(pushIndex).toBeGreaterThanOrEqual(0); + expect(paths.lastIndexOf('/api/v2/session/inputs/probe')).toBeGreaterThan(pushIndex); + expect(paths).not.toContain('/api/v2/execute'); + }); + + test('a payload with no by-ref inputs skips probe and push entirely', async () => { + const fake = fakeClient(); + const req = request(); + req.body.files = [{ name: 'inline.txt', content: 'inline' }]; + await makeBackend(fake).execute(req, sessionContext()); + expect(captured.filter((c) => c.path.startsWith('/api/v2/session/inputs'))).toHaveLength(0); + }); + + test('lock contention is retryable BUSY, never a cold one-shot', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake, { lockWaitMs: 10 }); + await acquireRuntimeSessionLock('rt_session_1', 60_000); + + /* A session-bound request depends on workspace state a fresh VM lacks — + * packages, files, database state — so answering from one would be + * silently wrong regardless of whether THIS payload carries refs. */ + for (const mode of ['affinity', 'strict'] as const) { + await expect( + backend.execute(request(), sessionContext({ runtimeSessionMode: mode })), + ).rejects.toThrow('busy'); + } + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + }); + + test('a fresh session VM returning a proxy 502 is recycled immediately', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + executeStatus = 502; + + await expect(backend.execute(request(), sessionContext())).rejects.toThrow(); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); + }); + + test('a reused VM returning a proxy 502 (failed auto-resume) is recycled', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + /* 502/503/504 is the AWS proxy reporting the VM unreachable (a suspended VM + * that failed to auto-resume), not the runner rejecting the request — so the + * dead VM must be torn down, unlike a runner 500. */ + executeStatus = 502; + await expect(backend.execute(request(), sessionContext())).rejects.toThrow(); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); + }); + + test('relaunches an idle-expired session instead of reusing the dead endpoint', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + /* Backdate last_seen past idle+suspended: AWS would have auto-terminated the + * VM, so the next request must relaunch rather than reuse the stale RUNNING + * endpoint (which would health-check-fail and 503 the first request). */ + const token = (await acquireRuntimeSessionLock('rt_session_1', 60_000)) as string; + const rec = await readRuntimeSessionRecord('rt_session_1'); + await writeRuntimeSessionRecord({ ...rec!, last_seen_at: 1 }, token); + const { releaseRuntimeSessionLock } = await import('../runtime-session/registry'); + await releaseRuntimeSessionLock('rt_session_1', token); + await backend.execute(request(), sessionContext()); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + test('two concurrent executions on one session serialize on the registry lock', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + + const [a, b] = await Promise.all([ + backend.execute(request(), sessionContext()), + backend.execute(request(), sessionContext()), + ]); + expect(a).toEqual(EXECUTE_RESPONSE); + expect(b).toEqual(EXECUTE_RESPONSE); + /* Serialized launch: exactly one VM created, reused by the other. */ + expect(fake.callsFor('runMicrovm')).toHaveLength(1); + }); + + test('a fenced post-execute record write fails instead of returning stale session state', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + stealSessionLockOnExecute = true; + + try { + await backend.execute(request(), sessionContext()); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('MICROVM_FENCED'); + } + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + test('a missing post-execute session record fences and terminates the known VM', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + onExecute = async () => { + await mock.del('rtsx:sess:rt_session_1'); + }; + + try { + await backend.execute(request(), sessionContext()); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('MICROVM_FENCED'); + } + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); + }); + + test('strict mode raises RUNTIME_SESSION_BUSY when the lock is held', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake, { lockWaitMs: 100 }); + const held = await acquireRuntimeSessionLock('rt_session_1', 60_000); + expect(held).not.toBeNull(); + + try { + await backend.execute(request(), sessionContext({ runtimeSessionMode: 'strict' })); + throw new Error('expected rejection'); + } catch (error) { + expect(error).toBeInstanceOf(SandboxBackendError); + expect((error as SandboxBackendError).code).toBe('RUNTIME_SESSION_BUSY'); + } + }); + + + test('stateless mode ignores a runtime session id', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext({ runtimeSessionMode: 'stateless' })); + const runArgs = fake.callsFor('runMicrovm')[0].args as { runHookPayload?: string }; + expect(runArgs.runHookPayload).toBeUndefined(); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); + }); + + test('session mode refuses an unpinned image version before launch', async () => { + const fake = fakeClient(); + await expect( + makeBackend(fake, { imageVersion: undefined }).execute(request(), sessionContext()), + ).rejects.toThrow('pinned LAMBDA_MICROVM_IMAGE_VERSION'); + expect(fake.callsFor('runMicrovm')).toHaveLength(0); + }); + + test('sends X-aws-proxy-port only when the port is not the 8080 default', async () => { + const fake = fakeClient(); + await makeBackend(fake, { port: 9090 }).execute(request(), sessionContext()); + const exec = captured.find((c) => c.path === '/api/v2/execute'); + /* Non-default port needs the routing header, or AWS sends traffic to 8080. */ + expect(exec?.headers['x-aws-proxy-port']).toBe('9090'); + + captured = []; + await makeBackend(fake, { port: 8080 }).execute(request(), sessionContext({ runtimeSessionId: 'rt_8080' })); + const exec8080 = captured.find((c) => c.path === '/api/v2/execute'); + expect(exec8080?.headers['x-aws-proxy-port']).toBeUndefined(); + }); + + test('a reused VM whose token mint returns not_found is torn down and the record dropped', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake); + await backend.execute(request(), sessionContext()); + /* The VM was evicted between calls: CreateMicrovmAuthToken now 404s. That + * escapes raw today; it must surface as MICROVM_UNHEALTHY so the dead VM is + * terminated and its record dropped, and the next call relaunches. */ + fake.failNext('createMicrovmAuthToken', new LambdaMicrovmApiError('not_found', 'CreateMicrovmAuthToken', 'gone')); + await expect(backend.execute(request(), sessionContext())).rejects.toBeInstanceOf(SandboxBackendError); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); + }); + + test('terminates a superseded (config-drifted) VM before relaunching', async () => { + const fake = fakeClient(); + await makeBackend(fake).execute(request(), sessionContext()); + const oldVmId = [...fake.vms.keys()][0]; + /* A deploy bumps the image version: the recorded VM no longer matches config, + * so it must be terminated (not left running/billing) before the replacement + * launches. */ + await makeBackend(fake, { imageVersion: '4' }).execute(request(), sessionContext()); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + const terminated = fake.callsFor('terminateMicrovm').map((c) => (c.args as { microvmId: string }).microvmId); + expect(terminated).toContain(oldVmId); + }); + + test('provider not-found while terminating a stale VM still permits relaunch', async () => { + const fake = fakeClient(); + await makeBackend(fake).execute(request(), sessionContext()); + fake.failNext( + 'terminateMicrovm', + new LambdaMicrovmApiError('not_found', 'TerminateMicrovm', 'already gone'), + ); + + await makeBackend(fake, { imageVersion: '4' }).execute(request(), sessionContext()); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + expect((await readRuntimeSessionRecord('rt_session_1'))?.image_version).toBe('4'); + }); + + test('launch-record failure after RunMicrovm terminates the untracked VM', async () => { + const fake = fakeClient(); + const redisWithScript = mock as unknown as { + writeRuntimeSessionRecordScript: (...args: string[]) => Promise; + }; + const originalWrite = redisWithScript.writeRuntimeSessionRecordScript.bind(mock); + let writes = 0; + redisWithScript.writeRuntimeSessionRecordScript = async (...args: string[]) => { + writes += 1; + if (writes === 2) throw new Error('Redis unavailable after launch'); + return originalWrite(...args); + }; + + await expect(makeBackend(fake).execute(request(), sessionContext())).rejects.toThrow( + 'Redis unavailable after launch', + ); + expect(fake.callsFor('runMicrovm')).toHaveLength(1); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + test('a tightened egress connector config makes an existing session non-reusable', async () => { + const fake = fakeClient(); + await makeBackend(fake).execute(request(), sessionContext()); + const oldVmId = [...fake.vms.keys()][0]; + /* Connectors apply only at RunMicrovm, so a hardened deploy that tightens + * egress must relaunch rather than keep serving on the old broader policy. */ + await makeBackend(fake, { + egressConnectorArns: ['arn:aws:lambda:us-east-2:1:network-connector:vpc-egress'], + }).execute(request(), sessionContext()); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + const terminated = fake.callsFor('terminateMicrovm').map((c) => (c.args as { microvmId: string }).microvmId); + expect(terminated).toContain(oldVmId); + }); +}); + +describe('LambdaMicrovmSandboxBackend auto-checkpoint', () => { + function sessionContext(overrides: Partial = {}): SandboxExecuteContext { + return context({ + runtimeSessionId: 'rt_ckpt_1', + runtimeSessionMode: 'affinity', + tenantId: 'tenant-a', + canonicalUserId: 'user-1', + ...overrides, + }); + } + const cfgOn: Partial = { checkpointsEnabled: true }; + + test('checkpoints the workspace after a session exec and records the pointer', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + const backend = makeBackend(fake, cfgOn, store); + + await backend.execute(request(), sessionContext()); + + const checkpoints = captured.filter((c) => c.path === '/api/v2/session/checkpoint'); + expect(checkpoints).toHaveLength(1); + /* The runner binds session mode from this header (hookless): without it the + * checkpoint/restore handlers 409 and state is lost across expiry. */ + expect(checkpoints[0].headers['x-runtime-session-id']).toBe('rt_ckpt_1'); + const stored = await store.get('rt_ckpt_1', 1_000_000); + expect(stored).not.toBeNull(); + try { + expect(await fsp.readFile(stored!.path, 'utf8')).toBe(checkpointBlob); + } finally { + await stored?.cleanup(); + } + const record = await readRuntimeSessionRecord('rt_ckpt_1'); + /* Key is an immutable, zero-padded per-session sequence. */ + expect(record?.workspace_checkpoint).toStartWith(checkpointPrefixFor('rt_ckpt_1')); + expect(record?.workspace_checkpoint).toEndWith('.tar.gz'); + expect(record?.checkpointed_at).toBeGreaterThan(0); + }); + + test('restores the prior checkpoint on retry when result finalization fails', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + await store.put('rt_ckpt_1', 7, Buffer.from('PRIOR_WORKSPACE')); + await store.commit('rt_ckpt_1', 7); + const checkpointKey = checkpointObjectKey('rt_ckpt_1', 7); + const seedToken = await acquireRuntimeSessionLock('rt_ckpt_1', 60_000); + await writeRuntimeSessionRecord({ + runtime_session_id: 'rt_ckpt_1', + tenant_id: 'tenant-a', + canonical_user_id: 'user-1', + state: 'TERMINATED', + generation: 1, + last_seen_at: 1, + workspace_checkpoint: checkpointKey, + }, seedToken as string); + await releaseRuntimeSessionLock('rt_ckpt_1', seedToken as string); + const backend = makeBackend(fake, cfgOn, store); + let finalizeAttempts = 0; + const sessionResultFinalizer = async (result: SandboxRawResponse): Promise => { + finalizeAttempts += 1; + if (finalizeAttempts === 1) { + throw new Error('egress gateway restore unavailable'); + } + return { ...result, session_id: 'restored_result' }; + }; + + await expect(backend.execute( + request(), + sessionContext({ sessionResultFinalizer }), + )).rejects.toThrow('egress gateway restore unavailable'); + + const firstAttemptPaths = captured.map((c) => c.path); + expect(firstAttemptPaths.filter((path) => path === '/api/v2/session/checkpoint')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + const recycled = await readRuntimeSessionRecord('rt_ckpt_1'); + expect(recycled?.state).toBe('TERMINATED'); + expect(recycled?.workspace_checkpoint).toBe(checkpointKey); + + const result = await backend.execute( + request(), + sessionContext({ sessionResultFinalizer }), + ); + expect(result.session_id).toBe('restored_result'); + expect(fake.callsFor('runMicrovm')).toHaveLength(2); + + const paths = captured.map((c) => c.path); + const executeIndexes = paths + .map((path, index) => path === '/api/v2/execute' ? index : -1) + .filter((index) => index >= 0); + const restoreIndexes = paths + .map((path, index) => path === '/api/v2/session/restore' ? index : -1) + .filter((index) => index >= 0); + expect(executeIndexes).toHaveLength(2); + expect(restoreIndexes).toHaveLength(2); + expect(restoreIndexes[1]).toBeGreaterThan(executeIndexes[0]); + expect(restoreIndexes[1]).toBeLessThan(executeIndexes[1]); + }); + + test('still terminates the mutated VM when the rollback registry read fails', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake, cfgOn, new MemoryCheckpointStore()); + const redisWithGet = mock as unknown as { + get(key: string): Promise; + }; + const originalGet = redisWithGet.get.bind(mock); + let failSessionRead = false; + redisWithGet.get = async (key: string): Promise => { + if (failSessionRead && key === 'rtsx:sess:rt_ckpt_1') { + throw new Error('registry read unavailable'); + } + return originalGet(key); + }; + const sessionResultFinalizer = async ( + _result: SandboxRawResponse, + ): Promise => { + failSessionRead = true; + throw new Error('egress gateway restore unavailable'); + }; + + try { + await expect(backend.execute( + request(), + sessionContext({ sessionResultFinalizer }), + )).rejects.toThrow('egress gateway restore unavailable'); + } finally { + failSessionRead = false; + } + + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + /* The registry could not be quarantined, so provider teardown—not Redis + * state—is what prevents reuse of the uncommitted workspace. */ + expect((await readRuntimeSessionRecord('rt_ckpt_1'))?.state).toBe('RUNNING'); + }); + + test('still terminates the mutated VM when the rollback quarantine write fails', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake, cfgOn, new MemoryCheckpointStore()); + const redisWithScript = mock as unknown as { + writeRuntimeSessionRecordScript(...args: string[]): Promise; + }; + const originalWrite = redisWithScript.writeRuntimeSessionRecordScript.bind(mock); + let failQuarantineWrite = false; + redisWithScript.writeRuntimeSessionRecordScript = async ( + ...args: string[] + ): Promise => { + if (failQuarantineWrite) throw new Error('registry write unavailable'); + return originalWrite(...args); + }; + const sessionResultFinalizer = async ( + _result: SandboxRawResponse, + ): Promise => { + failQuarantineWrite = true; + throw new Error('egress gateway restore unavailable'); + }; + + try { + await expect(backend.execute( + request(), + sessionContext({ sessionResultFinalizer }), + )).rejects.toThrow('egress gateway restore unavailable'); + } finally { + failQuarantineWrite = false; + } + + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect((await readRuntimeSessionRecord('rt_ckpt_1'))?.state).toBe('RUNNING'); + }); + + test('a relaunched VM restores the checkpoint before the first exec', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_ckpt_1', 1000, Buffer.from('PRIOR_WORKSPACE')); + await store.commit('rt_ckpt_1', 1000); + /* Seed a terminated prior session so findOrLaunch relaunches. */ + const seedToken = await acquireRuntimeSessionLock('rt_ckpt_1', 60_000); + await writeRuntimeSessionRecord({ + runtime_session_id: 'rt_ckpt_1', tenant_id: 'tenant-a', canonical_user_id: 'user-1', + state: 'TERMINATED', generation: 3, last_seen_at: 1, workspace_checkpoint: checkpointObjectKey('rt_ckpt_1', 1000), + }, seedToken as string); + const { releaseRuntimeSessionLock } = await import('../runtime-session/registry'); + await releaseRuntimeSessionLock('rt_ckpt_1', seedToken as string); + + const fake = fakeClient(); + const backend = makeBackend(fake, cfgOn, store); + const result = await backend.execute(request(), sessionContext()); + expect(result).toEqual(EXECUTE_RESPONSE); + + const paths = captured.map((c) => c.path); + /* restore precedes execute on the fresh VM. */ + expect(paths.indexOf('/api/v2/session/restore')).toBeGreaterThanOrEqual(0); + expect(paths.indexOf('/api/v2/session/restore')).toBeLessThan(paths.indexOf('/api/v2/execute')); + const restoreReq = captured.find((c) => c.path === '/api/v2/session/restore'); + expect(restoreReq?.headers['x-runtime-session-id']).toBe('rt_ckpt_1'); + expect(recordDuringRestore?.state).toBe('PENDING'); + expect(recordDuringRestore?.workspace_checkpoint).toBe( + checkpointObjectKey('rt_ckpt_1', 1000), + ); + expect(fake.callsFor('runMicrovm')).toHaveLength(1); + }); + + test('reuse (warm VM) does not restore — no prior expiry', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + const backend = makeBackend(fake, cfgOn, store); + + await backend.execute(request(), sessionContext()); + captured = []; + await backend.execute(request(), sessionContext()); + + expect(captured.filter((c) => c.path === '/api/v2/session/restore')).toHaveLength(0); + expect(fake.callsFor('runMicrovm')).toHaveLength(1); + }); + + test('disabled checkpoints skip both checkpoint and restore', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + const backend = makeBackend(fake, { checkpointsEnabled: false }, store); + await backend.execute(request(), sessionContext()); + /* File delivery is independent of checkpointing — only the checkpoint and + * restore legs must be skipped. */ + expect(captured.filter((c) => + c.path === '/api/v2/session/checkpoint' || c.path === '/api/v2/session/restore', + )).toHaveLength(0); + expect(store.objects.size).toBe(0); + }); + + test('skips a checkpoint unless the budget covers all four bounded I/O legs', async () => { + const originalNow = Date.now; + let now = 1_000_000; + Date.now = () => now; + onExecute = () => { + /* Leaves 45ms: more than the old 10 + 3*10 guard, but less than the + * actual token wait + GET + list + put + commit worst case (50ms). */ + now += 55; + }; + try { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + const backend = makeBackend(fake, { + checkpointsEnabled: true, + jobTimeoutMs: 100, + launchTimeoutMs: 10, + checkpoint: { + port: 8080, + authTokenTtlSeconds: 300, + maxBytes: 512 * 1024 * 1024, + timeoutMs: 10, + }, + }, store); + + expect(await backend.execute(request(), sessionContext())).toEqual(EXECUTE_RESPONSE); + expect(captured.filter((c) => c.path === '/api/v2/session/checkpoint')).toHaveLength(0); + expect(store.objects.size).toBe(0); + } finally { + Date.now = originalNow; + } + }); + + test('production defaults leave enough budget for a fresh session checkpoint', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + const backend = makeBackend(fake, { + checkpointsEnabled: true, + jobTimeoutMs: 300_000, + launchTimeoutMs: 60_000, + checkpoint: { + port: 8080, + authTokenTtlSeconds: 300, + maxBytes: 512 * 1024 * 1024, + timeoutMs: 60_000, + }, + }, store); + + expect(await backend.execute(request(), sessionContext())).toEqual(EXECUTE_RESPONSE); + expect(captured.filter((c) => c.path === '/api/v2/session/checkpoint')).toHaveLength(1); + expect(store.objects.size).toBe(1); + }); + + test('uses the worker deadline so pre-backend setup consumes checkpoint budget', async () => { + const originalNow = Date.now; + const now = 1_000_000; + Date.now = () => now; + try { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + const backend = makeBackend(fake, { + checkpointsEnabled: true, + /* A backend-local 100ms clock would allow the 50ms pipeline. The + * worker deadline has only 49ms left after earlier egress/setup work. */ + jobTimeoutMs: 100, + launchTimeoutMs: 10, + checkpoint: { + port: 8080, + authTokenTtlSeconds: 300, + maxBytes: 512 * 1024 * 1024, + timeoutMs: 10, + }, + }, store); + + expect(await backend.execute( + request(), + sessionContext({ deadlineAtMs: now + 49 }), + )).toEqual(EXECUTE_RESPONSE); + expect(captured.filter((c) => c.path === '/api/v2/session/checkpoint')).toHaveLength(0); + expect(store.objects.size).toBe(0); + } finally { + Date.now = originalNow; + } + }); + + test('a checkpoint FETCH failure fails closed instead of running on an empty workspace', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + store.get = () => Promise.reject(new Error('S3 down')); + const backend = makeBackend(fake, cfgOn, store); + + /* Running anyway used to let the post-run checkpoint prune the last good + * snapshot — a transient S3 blip becoming permanent data loss. */ + await expect(backend.execute(request(), sessionContext())).rejects.toThrow( + 'refusing to run against an empty workspace', + ); + expect(captured.filter((c) => c.path === '/api/v2/execute')).toHaveLength(0); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + expect(await readRuntimeSessionRecord('rt_ckpt_1')).toBeNull(); + }); + + test('a failed checkpoint is non-fatal — the exec still succeeds', async () => { + const fake = fakeClient(); + const failing: MemoryCheckpointStore = new MemoryCheckpointStore(); + failing.put = () => Promise.reject(new Error('S3 down')); + const backend = makeBackend(fake, cfgOn, failing); + const result = await backend.execute(request(), sessionContext()); + expect(result).toEqual(EXECUTE_RESPONSE); + const record = await readRuntimeSessionRecord('rt_ckpt_1'); + expect(record?.state).toBe('RUNNING'); + expect(record?.workspace_checkpoint).toBeUndefined(); + }); + + test('a failed first reseed lookup cannot strand the Redis sequence below retained objects', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + await store.put('rt_ckpt_1', 100, Buffer.from('PRIOR_WORKSPACE')); + await store.commit('rt_ckpt_1', 100); + const latestSequence = store.latestSequence.bind(store); + let lookups = 0; + store.latestSequence = async (runtimeSessionId) => { + lookups += 1; + if (lookups === 1) throw new Error('transient list failure'); + return latestSequence(runtimeSessionId); + }; + const backend = makeBackend(fake, cfgOn, store); + + expect(await backend.execute(request(), sessionContext())).toEqual(EXECUTE_RESPONSE); + expect(await backend.execute(request(), sessionContext())).toEqual(EXECUTE_RESPONSE); + + const record = await readRuntimeSessionRecord('rt_ckpt_1'); + expect(record?.workspace_checkpoint).toBe(checkpointObjectKey('rt_ckpt_1', 101)); + expect(store.objects.has(checkpointObjectKey('rt_ckpt_1', 101))).toBe(true); + }); + + test('a commit-marker failure keeps the previous durable recovery point', async () => { + const fake = fakeClient(); + const store = new MemoryCheckpointStore(); + await store.put('rt_ckpt_1', 10, Buffer.from('PRIOR_WORKSPACE')); + await store.commit('rt_ckpt_1', 10); + const commit = store.commit.bind(store); + store.commit = async (runtimeSessionId, sequence) => { + if (sequence === 11) throw new Error('marker write failed'); + await commit(runtimeSessionId, sequence); + }; + let pruneCalls = 0; + const prune = store.pruneOlderThan.bind(store); + store.pruneOlderThan = async (runtimeSessionId, sequence) => { + pruneCalls += 1; + await prune(runtimeSessionId, sequence); + }; + const backend = makeBackend(fake, cfgOn, store); + + expect(await backend.execute(request(), sessionContext())).toEqual(EXECUTE_RESPONSE); + expect((await readRuntimeSessionRecord('rt_ckpt_1'))?.workspace_checkpoint).toBe( + checkpointObjectKey('rt_ckpt_1', 11), + ); + expect(pruneCalls).toBe(0); + + const durable = await store.get('rt_ckpt_1', 1_000_000); + try { + expect(await fsp.readFile(durable!.path, 'utf8')).toBe('PRIOR_WORKSPACE'); + } finally { + await durable?.cleanup(); + } + }); +}); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts new file mode 100644 index 0000000..8dbfd87 --- /dev/null +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -0,0 +1,1296 @@ +import axios from 'axios'; +import { nanoid } from 'nanoid'; +import * as fs from 'fs'; +import type { LambdaMicrovmClient, MicrovmAuthToken, MicrovmDescription, MicrovmIdlePolicy } from '../runtime-session/lambda-client'; +import type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +import type { RuntimeSessionRecord } from '../runtime-session/registry'; +import type { CheckpointConfig } from '../runtime-session/checkpoint'; +import type { CheckpointStore } from '../runtime-session/checkpoint-store'; +import { LambdaMicrovmApiError, microvmPortHeaders } from '../runtime-session/lambda-client'; +import { MicrovmOpThrottledError, acquireOpBudget, poisonOpBucket } from '../runtime-session/throttle'; +import { checkpointSession, probeInputs, pushInputs, restoreSession } from '../runtime-session/checkpoint'; +import { + SESSION_INPUTS_MAX_COUNT, + SessionFilesError, + buildInputBatch, + sessionFileRefs, +} from '../runtime-session/files'; +import { + RUNTIME_SESSION_LOCK_TTL_MS, + allocateRuntimeSessionGeneration, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + removeRuntimeSession, + renewRuntimeSessionLock, + waitForRuntimeSessionLock, + writeRuntimeSessionRecord, +} from '../runtime-session/registry'; +import { startRuntimeSessionLockHeartbeat } from '../runtime-session/lock-heartbeat'; +import { + microvmLaunches, + microvmLaunchDuration, + microvmTerminations, + microvmThrottleEvents, + runtimeSessionLockContention, +} from '../metrics'; +import { injectTraceHeaders, withSpan } from '../telemetry'; +import { SandboxBackendError } from './types'; +import { Jobs } from '../enum'; +import { checkpointPipelineBudgetMs } from '../config'; +import logger from '../logger'; + +/** Header that opts a proxied /execute into the runner's persistent session + * workspace (see api/src/session-workspace.ts). Session mode is delivered + * per-request, not via a /run lifecycle hook — Lambda's build hooks require + * the snapshot-compatible base container image to route, and enabling any + * runtime hook forces the /ready build hook (which never reaches a stock + * container's listener), so hookless + per-request keeps image builds sound. */ +const RUNTIME_SESSION_ID_HEADER = 'X-Runtime-Session-Id'; + +export interface LambdaMicrovmBackendConfig { + imageArn: string; + imageVersion?: string; + executionRoleArn?: string; + logGroup?: string; + ingressConnectorArns?: string[]; + egressConnectorArns?: string[]; + port: number; + maxDurationSeconds: number; + authTokenTtlSeconds: number; + launchTimeoutMs: number; + healthTimeoutMs: number; + launchTps: number; + tokenTps: number; + jobTimeoutMs: number; + /* Session-mode (find-or-launch) tuning. */ + idleSeconds: number; + suspendedSeconds: number; + lockWaitMs: number; + /* Auto-checkpoint. When disabled, session VMs still reuse a warm workspace + * but expiry recovery falls back to file refs (no cross-VM restore). */ + checkpointsEnabled: boolean; + checkpoint: CheckpointConfig; +} + +/** Order-independent fingerprint of every immutable launch/security input. */ +export function runtimeSessionLaunchFingerprint(config: LambdaMicrovmBackendConfig): string { + const ingress = [...(config.ingressConnectorArns ?? [])].sort(); + const egress = [...(config.egressConnectorArns ?? [])].sort(); + return JSON.stringify({ + imageArn: config.imageArn, + imageVersion: config.imageVersion, + executionRoleArn: config.executionRoleArn ?? '', + logGroup: config.logGroup ?? '', + ingress, + egress, + port: config.port, + maximumDurationSeconds: config.maxDurationSeconds, + idlePolicy: { + maxIdleSeconds: config.idleSeconds, + suspendedSeconds: config.suspendedSeconds, + autoResume: true, + }, + }); +} + +interface LambdaMicrovmBackendDeps { + clientFactory: () => Promise; + config: LambdaMicrovmBackendConfig; + pollIntervalMs?: number; + /** Injected in session+checkpoint mode; undefined disables checkpoints. */ + checkpointStore?: CheckpointStore; +} + +const sleep = (ms: number, signal?: AbortSignal): Promise => new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(signal.reason instanceof Error ? signal.reason : new Error('Operation aborted')); + return; + } + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + const onAbort = (): void => { + clearTimeout(timer); + reject(signal?.reason instanceof Error ? signal.reason : new Error('Operation aborted')); + }; + signal?.addEventListener('abort', onAbort, { once: true }); +}); + +/** AWS returns the endpoint as a bare host; docs samples do `https://${endpoint}`. */ +export function normalizeMicrovmEndpoint(endpoint: string): string { + if (endpoint.startsWith('http://') || endpoint.startsWith('https://')) { + return endpoint.replace(/\/+$/, ''); + } + return `https://${endpoint.replace(/\/+$/, '')}`; +} + +interface LaunchOptions { + clientToken: string; + idlePolicy?: MicrovmIdlePolicy; + maxDurationSeconds: number; +} + +/** + * Lambda MicroVM backend. Two modes, chosen by the runtime session context: + * + * - **stateless** (no runtime session): one VM per execution — run, execute, + * terminate. Correct and simple; the default. + * - **session** (affinity/strict): find-or-launch one warm VM per + * `runtime_session_id` via the registry, stamp that id on every proxied + * /execute (the header that activates the runner's persistent workspace), + * and reuse the VM across calls. + * AWS `idlePolicy` auto-suspends the VM when idle and auto-resumes it on the + * next request, so there is no explicit resume in the execute path. + */ +export class LambdaMicrovmSandboxBackend implements SandboxBackend { + readonly name = 'lambda-microvm' as const; + private clientPromise: Promise | undefined; + private readonly config: LambdaMicrovmBackendConfig; + private readonly clientFactory: () => Promise; + private readonly pollIntervalMs: number; + private readonly checkpointStore: CheckpointStore | undefined; + + constructor(deps: LambdaMicrovmBackendDeps) { + this.clientFactory = deps.clientFactory; + this.config = deps.config; + this.pollIntervalMs = deps.pollIntervalMs ?? 500; + this.checkpointStore = deps.checkpointStore; + } + + private checkpointsActive(): boolean { + return this.config.checkpointsEnabled && this.checkpointStore !== undefined; + } + + private client(): Promise { + this.clientPromise ??= this.clientFactory(); + return this.clientPromise; + } + + async execute(req: SandboxTransportRequest, ctx: SandboxExecuteContext): Promise { + /* Production supplies the worker-owned deadline, captured before egress + * grant and request setup. Keep a backend-entry fallback for direct/test + * callers that do not own a worker timer. Capture it before lazy client + * initialization so that setup is counted too. */ + const deadlineAtMs = ctx.deadlineAtMs ?? Date.now() + this.config.jobTimeoutMs; + const client = await this.client(); + if (ctx.runtimeSessionId && ctx.runtimeSessionMode !== 'stateless') { + return this.executeSession(client, req, ctx, ctx.runtimeSessionId, deadlineAtMs); + } + return this.executeStateless(client, req, ctx); + } + + private async executeStateless( + client: LambdaMicrovmClient, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + ): Promise { + /* One-shots self-cap their lifetime near the job timeout so a crashed + * worker cannot leak an 8h VM. */ + const maxDurationSeconds = Math.min( + this.config.maxDurationSeconds, + Math.ceil(this.config.jobTimeoutMs / 1_000) + 120, + ); + const vm = await this.launch(client, ctx, { + clientToken: ctx.executionId !== '' ? `exec-${ctx.executionId}` : `exec-${nanoid()}`, + maxDurationSeconds, + }); + let terminateReason = 'stateless'; + try { + await this.waitForRunnerReady( + client, + vm.microvmId, + normalizeMicrovmEndpoint(vm.endpoint ?? ''), + ctx, + ); + /* Stateless one-shots take by-reference inputs too, and their guest is + * just as unable to pull them. Same cache, same priming path. */ + await this.deliverInputs(client, vm, req, ctx); + return await this.proxyExecute(client, vm, req, ctx, undefined, true); + } catch (error) { + terminateReason = ctx.signal.aborted ? 'timeout' : 'error'; + throw error; + } finally { + await this.terminate(client, vm.microvmId, terminateReason); + } + } + + private async executeSession( + client: LambdaMicrovmClient, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + runtimeSessionId: string, + deadlineAtMs: number, + ): Promise { + const lockToken = await waitForRuntimeSessionLock(runtimeSessionId, { + waitMs: this.config.lockWaitMs, + signal: ctx.signal, + }); + if (!lockToken) { + runtimeSessionLockContention.inc({ mode: ctx.runtimeSessionMode }); + /* A session-bound request depends on that session's workspace — files, + * installed packages, database state built by earlier turns. Running it + * on a cold one-shot would silently answer from an environment the + * caller never asked for, so contention is a retryable BUSY in every + * session mode. (A lossy fallback, if ever wanted, belongs behind its + * own explicit mode rather than as a silent default.) */ + throw new SandboxBackendError( + 'RUNTIME_SESSION_BUSY', + `Runtime session ${runtimeSessionId} is busy`, + ); + } + + + const fence = new AbortController(); + const heartbeat = startRuntimeSessionLockHeartbeat({ + renew: () => renewRuntimeSessionLock(runtimeSessionId, lockToken), + fence, + ttlMs: RUNTIME_SESSION_LOCK_TTL_MS, + }); + const sessionCtx: SandboxExecuteContext = { + ...ctx, + signal: AbortSignal.any([ctx.signal, fence.signal]), + }; + let sessionVm: MicrovmDescription | undefined; + + try { + const existing = await readRuntimeSessionRecord(runtimeSessionId); + const { vm } = await this.findOrLaunchSession( + client, + sessionCtx, + runtimeSessionId, + existing, + lockToken, + ); + sessionVm = vm; + await this.deliverInputs(client, vm, req, sessionCtx, async () => { + /* A session VM whose delivery transport failed is not trustworthy for + * reuse: recycle so the next call relaunches and restores. */ + const terminated = await this.terminate(client, vm.microvmId, 'error'); + if (terminated) { + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + } + }); + const result = await this.executeOnSessionVm( + client, + vm, + req, + sessionCtx, + runtimeSessionId, + lockToken, + fence.signal, + ); + let finalizedResult: SandboxRawResponse; + try { + /* Gateway result restoration is part of the session transaction: do + * not checkpoint or advertise the mutated workspace until it succeeds. + * If it fails, quarantine/recycle this VM so an at-least-once retry + * relaunches from the preceding durable checkpoint. */ + finalizedResult = sessionCtx.sessionResultFinalizer + ? await sessionCtx.sessionResultFinalizer(result) + : result; + } catch (error) { + await this.recycleUncommittedSessionVm( + client, + vm, + runtimeSessionId, + lockToken, + 'result_finalization_failed', + ); + throw error; + } + /* Re-read the record findOrLaunch settled on (freshly written on + * launch, or the reused one) and only bump its liveness — preserves + * generation, deadline, and image fields. */ + const now = Date.now(); + /* The post-run checkpoint is an optional cache write. Skip it when the job + * budget won't fit a full checkpoint, so a run that already succeeded + * isn't timed out at the router by the checkpoint's latency — the next + * relaunch restores the prior checkpoint, one exec staler. */ + const remainingBudgetMs = deadlineAtMs - now; + /* Reserve the WHOLE checkpoint path, not just one transfer timeout: + * token-budget wait + guest GET + object-store list + object upload + + * durable marker. Metadata calls use their tighter cap; the two archive + * transfers receive the configured checkpoint timeout. */ + const worstCaseCheckpointMs = checkpointPipelineBudgetMs( + this.config.launchTimeoutMs, + this.config.checkpoint.timeoutMs, + ); + const canCheckpoint = !sessionCtx.signal.aborted && remainingBudgetMs > worstCaseCheckpointMs; + const settled = await readRuntimeSessionRecord(runtimeSessionId); + if (!settled) { + /* Returning success here would strand the known live VM without a + * registry record: no later caller could reuse or clean it up. Treat + * disappearance as fencing so the catch path terminates this VM. */ + throw new SandboxBackendError( + 'MICROVM_FENCED', + `Runtime session record disappeared for ${runtimeSessionId} after execute`, + ); + } + const nextRecord = canCheckpoint + ? await this.checkpointUnderLock( + client, + settled, + runtimeSessionId, + now, + lockToken, + sessionCtx.signal, + ) + : { ...settled, state: 'RUNNING' as const, last_seen_at: now }; + const persisted = await writeRuntimeSessionRecord(nextRecord, lockToken); + if (!persisted) { + throw new SandboxBackendError('MICROVM_FENCED', `Lost session lock for ${runtimeSessionId} after execute`); + } + return finalizedResult; + } catch (error) { + /* Fencing can happen after proxyExecute returns — during checkpoint, + * the final registry read, or the final liveness write. At that point a + * newer holder may already be using the recorded VM. Kill it so the new + * holder fails fast and restores, rather than letting two workers touch + * one persistent workspace concurrently. The stale lock token cannot + * safely remove the new holder's record, so termination is the only + * mutation here. */ + const fenced = + fence.signal.aborted || + (error instanceof SandboxBackendError && error.code === 'MICROVM_FENCED'); + if (fenced && sessionVm) { + await this.terminate(client, sessionVm.microvmId, 'fenced'); + } + throw error; + } finally { + heartbeat.stop(); + await releaseRuntimeSessionLock(runtimeSessionId, lockToken); + } + } + + /** + * Proxies the execute to a session VM and, on a failure that means the VM + * must not be reused, terminates it and makes the registry record + * non-reusable so the next call relaunches + restores rather than reusing a + * dead-or-dirty VM: + * - abort (JOB_TIMEOUT): the runner keeps NsJail running until the child + * exits even after the socket closes, so a later request reusing this VM + * could mutate the workspace concurrently with the timed-out run. + * - VM unreachable (health/connection failure, e.g. idlePolicy auto-terminated + * a suspended VM): the RUNNING record would otherwise keep pointing at a + * dead VM until the hard deadline, and every request would reuse it. + * A plain non-200 from a live runner (`Error from sandbox`) leaves the warm VM + * and its record intact — the VM is healthy, only the request failed. + */ + private async executeOnSessionVm( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + runtimeSessionId: string, + lockToken: string, + fenceSignal?: AbortSignal, + ): Promise { + try { + /* Fresh sessions passed readiness before their PENDING record became + * RUNNING. Reused sessions intentionally let probe/execute trigger + * auto-resume under the full job budget. */ + return await this.proxyExecute(client, vm, req, ctx, runtimeSessionId, true); + } catch (error) { + /* Recycle the VM ONLY on positive evidence it's unreachable or dirty: + * - abort: the runner keeps NsJail running after the socket closes, so a + * reuse could mutate the workspace concurrently. + * - a transport-level axios failure (no `.response`): the execute couldn't + * reach the VM (connection refused/timeout). + * - a failed health check: assertHealthy wraps the connection/timeout/non-200 + * into MICROVM_UNHEALTHY, so it isn't a top-level AxiosError. + * Everything else keeps the warm VM: a non-2xx sandbox response (AxiosError + * WITH `.response` — the VM is alive, only the request failed) and a + * pre-request control-plane failure like a throttled CreateMicrovmAuthToken + * (not an axios error at all) — the VM was never touched. + * Exception: a proxy 5xx (502/503/504) is the AWS gateway reporting the + * VM as unreachable — typically a suspended VM that failed to auto-resume, + * but it can also happen after a fresh VM passed health and then died. + * That has `.response` (so it's not a transport failure) but means the VM + * is dead, not that the runner rejected the request; recycle it, else + * every later call keeps reusing the dead VM until idle expiry. */ + const status = axios.isAxiosError(error) ? error.response?.status ?? 0 : 0; + /* Fenced: another worker owns the lock and will reuse this VM. Aborting + * our HTTP request does NOT stop the runner's NsJail child (it runs to + * completion after the socket closes), so leaving the VM alive would let + * two executions mutate one persistent workspace concurrently. Terminate + * it: the new holder's call fails fast on the dead endpoint and + * relaunches + restores from the last checkpoint. A surfaced error and a + * relaunch beat silent workspace corruption. The registry write is a + * no-op under our stale token, which is fine — the reuse path already + * recycles records pointing at dead VMs. */ + if (fenceSignal?.aborted === true) { + throw new SandboxBackendError( + 'MICROVM_FENCED', + `Lost session lock for ${runtimeSessionId} during execute`, + error, + ); + } + const transportFailure = axios.isAxiosError(error) && error.response == null; + const unhealthy = error instanceof SandboxBackendError && error.code === 'MICROVM_UNHEALTHY'; + const gatewayUnreachable = status >= 502 && status <= 504; + const sessionWorkspaceDirty = + axios.isAxiosError(error) && + error.response?.status === 409 && + (error.response.data as { error?: unknown } | undefined)?.error === 'session_workspace_dirty'; + if (sessionWorkspaceDirty) { + await this.recycleUncommittedSessionVm( + client, + vm, + runtimeSessionId, + lockToken, + 'workspace_dirty', + ); + throw new SandboxBackendError( + 'MICROVM_UNHEALTHY', + `Runtime session ${runtimeSessionId} workspace was dirty and has been recycled`, + error, + ); + } + if ( + ctx.signal.aborted || + transportFailure || + unhealthy || + gatewayUnreachable + ) { + const terminated = await this.terminate( + client, + vm.microvmId, + ctx.signal.aborted ? 'timeout' : 'error', + ); + if (terminated) { + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + } + } + throw error; + } + } + + /** + * Makes a session VM non-reusable before attempting provider cleanup. The + * prior checkpoint pointer remains on the terminal record, so a retry can + * relaunch directly from the last committed workspace even if the durable + * marker write for that older checkpoint had failed. + */ + private async recycleUncommittedSessionVm( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + runtimeSessionId: string, + lockToken: string, + reason: 'result_finalization_failed' | 'workspace_dirty', + ): Promise { + let record: RuntimeSessionRecord | null = null; + let ownsRecordedVm = false; + let quarantined = false; + try { + record = await readRuntimeSessionRecord(runtimeSessionId); + ownsRecordedVm = record?.microvm_id === vm.microvmId; + if (record && ownsRecordedVm) { + quarantined = await writeRuntimeSessionRecord( + { ...record, state: 'TERMINATING', last_error: reason }, + lockToken, + ); + if (!quarantined) { + logger.warn('Lost session lock while quarantining a mutated MicroVM', { + runtimeSessionId, + microvmId: vm.microvmId, + reason, + }); + } + } + } catch (error) { + /* Redis can fail exactly when the workspace most needs quarantining. + * Registry state is useful for making retries efficient, but provider + * termination is the actual safety boundary: never let a failed read or + * fenced write skip teardown of a VM whose uncommitted workspace was + * already mutated. */ + logger.error('Failed to quarantine a mutated MicroVM in the session registry', { + runtimeSessionId, + microvmId: vm.microvmId, + reason, + error, + }); + } + + const terminated = await this.terminate(client, vm.microvmId, reason); + if (!terminated || !record || !ownsRecordedVm || !quarantined) return; + + try { + const terminal = await writeRuntimeSessionRecord( + { + ...record, + microvm_id: undefined, + endpoint: undefined, + state: 'TERMINATED', + last_seen_at: Date.now(), + last_error: reason, + }, + lockToken, + ); + if (!terminal) { + logger.warn('Lost session lock after recycling a mutated MicroVM', { + runtimeSessionId, + microvmId: vm.microvmId, + reason, + }); + } + } catch (error) { + /* The provider resource is already gone and the earlier TERMINATING + * record remains non-reusable with its prior checkpoint pointer intact. + * Preserve the primary execution/finalization error. */ + logger.error('Failed to mark a recycled MicroVM terminal in the session registry', { + runtimeSessionId, + microvmId: vm.microvmId, + reason, + error, + }); + } + } + + /** Order-independent fingerprint of every immutable launch/security input. */ + private launchFingerprint(): string { + return runtimeSessionLaunchFingerprint(this.config); + } + + private async findOrLaunchSession( + client: LambdaMicrovmClient, + ctx: SandboxExecuteContext, + runtimeSessionId: string, + record: RuntimeSessionRecord | null, + lockToken: string, + ): Promise<{ vm: MicrovmDescription; restored: boolean }> { + if (!this.config.imageVersion) { + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + 'Stateful runtime sessions require a pinned LAMBDA_MICROVM_IMAGE_VERSION', + ); + } + const deadlineHeadroomMs = this.config.jobTimeoutMs + 30_000; + /* A record whose image/version/port no longer match the current config was + * launched by an older deploy — relaunch on the current config rather than + * reuse it (a changed port would otherwise health-check the wrong port and + * fail as UNHEALTHY instead of cleanly relaunching). */ + const configMatches = record + && record.image_arn === this.config.imageArn + && record.image_version === this.config.imageVersion + && record.port === this.config.port + && record.launch_fingerprint === this.launchFingerprint(); + /* Past idle+suspended, AWS auto-terminates the suspended VM while the record + * still reads RUNNING until the 8h hard deadline. Treat that as non-reusable + * so the first request after idle expiry relaunches + restores, instead of + * reusing a dead endpoint, failing the health check, and returning 503. */ + const idleTerminationMs = (this.config.idleSeconds + this.config.suspendedSeconds) * 1_000; + const likelyIdleTerminated = record != null && Date.now() - record.last_seen_at > idleTerminationMs; + const reusable = record + && record.state === 'RUNNING' + && record.microvm_id + && record.endpoint + && configMatches + && !likelyIdleTerminated + && (record.hard_deadline_at == null || record.hard_deadline_at - Date.now() > deadlineHeadroomMs); + if (reusable && record) { + /* Reuse the warm VM. If AWS auto-suspended it, the proxy request + * transparently auto-resumes it (idlePolicy.autoResume). */ + return { + vm: { microvmId: record.microvm_id as string, state: 'RUNNING', endpoint: record.endpoint }, + restored: false, + }; + } + + /* We're relaunching. Always terminate a recorded VM before replacing it. + * `last_seen_at` is only advanced after a successful execute; a failed + * delivery/proxy request can nevertheless auto-resume and touch the VM. + * Therefore "past idle+suspended" is enough to reject reuse, but NOT proof + * AWS already killed it. Skipping termination on that assumption leaks a + * live resumed VM after the record is overwritten. Provider not-found is + * already treated as successful cleanup, so the safe path is unconditional. */ + if (record?.microvm_id) { + const terminated = await this.terminate(client, record.microvm_id, 'superseded'); + if (!terminated) { + throw new SandboxBackendError( + 'MICROVM_UNHEALTHY', + `Could not terminate non-reusable MicroVM ${record.microvm_id}`, + ); + } + } + + /* RunMicrovm is idempotent by clientToken, but a worker can die after AWS + * accepts the call and before its response (and MicroVM id) reaches us. + * Persist the complete launch intent first. A successor on the same config + * replays that generation/token and recovers the already-running VM instead + * of allocating a new generation and leaving the first VM billable. */ + const canReplayPendingLaunch = Boolean( + record + && record.state === 'PENDING' + && !record.microvm_id + && configMatches + && Number.isSafeInteger(record.generation) + && record.generation > 0 + && record.launched_at != null + && record.hard_deadline_at != null + && record.hard_deadline_at - Date.now() > deadlineHeadroomMs, + ); + const generation = canReplayPendingLaunch && record + ? record.generation + : await allocateRuntimeSessionGeneration(runtimeSessionId); + const launchedAt = canReplayPendingLaunch && record?.launched_at != null + ? record.launched_at + : Date.now(); + const hardDeadlineAt = canReplayPendingLaunch && record?.hard_deadline_at != null + ? record.hard_deadline_at + : launchedAt + this.config.maxDurationSeconds * 1_000 - 60_000; + const launchIntent: RuntimeSessionRecord = { + runtime_session_id: runtimeSessionId, + tenant_id: ctx.tenantId ?? '', + canonical_user_id: ctx.canonicalUserId ?? '', + port: this.config.port, + image_arn: this.config.imageArn, + image_version: this.config.imageVersion, + launch_fingerprint: this.launchFingerprint(), + state: 'PENDING', + generation, + launched_at: launchedAt, + last_seen_at: Date.now(), + hard_deadline_at: hardDeadlineAt, + workspace_checkpoint: record?.workspace_checkpoint, + checkpointed_at: record?.checkpointed_at, + }; + const pendingOk = await writeRuntimeSessionRecord(launchIntent, lockToken); + if (!pendingOk) { + throw new SandboxBackendError('MICROVM_FENCED', `Lost session lock for ${runtimeSessionId} before launch`); + } + + const vm = await this.launch(client, ctx, { + clientToken: `sess-${runtimeSessionId}-${generation}`, + idlePolicy: { + maxIdleSeconds: this.config.idleSeconds, + suspendedSeconds: this.config.suspendedSeconds, + autoResume: true, + }, + maxDurationSeconds: this.config.maxDurationSeconds, + }); + + const launchedRecord: RuntimeSessionRecord = { + ...launchIntent, + microvm_id: vm.microvmId, + endpoint: vm.endpoint, + image_arn: vm.imageArn ?? this.config.imageArn, + image_version: vm.imageVersion ?? this.config.imageVersion, + last_seen_at: Date.now(), + }; + + let restored = false; + try { + /* Publish the launched VM as PENDING before readiness/restore. A worker + * crash leaves enough information for the successor to terminate it, + * while never advertising a partial workspace as reusable. */ + const tracked = await writeRuntimeSessionRecord(launchedRecord, lockToken); + if (!tracked) { + throw new SandboxBackendError( + 'MICROVM_FENCED', + `Lost session lock for ${runtimeSessionId} after launch`, + ); + } + + const endpointBase = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + await this.waitForRunnerReady(client, vm.microvmId, endpointBase, ctx); + + /* Fresh VM: restore the predecessor's workspace before the first execute + * so an 8h rollover / eviction is invisible. Attempt even after Redis + * record loss because the object-store key is deterministic. */ + if (this.checkpointStore && this.checkpointsActive()) { + const restoreResult = await restoreSession({ + mintToken: (microvmId) => this.mintAuthToken(client, microvmId, ctx.signal), + store: this.checkpointStore, + runtimeSessionId, + microvmId: vm.microvmId, + endpointBase, + config: this.config.checkpoint, + signal: ctx.signal, + checkpointKey: record?.workspace_checkpoint, + }); + if (restoreResult === 'push_failed' || restoreResult === 'fetch_failed') { + throw new SandboxBackendError( + 'MICROVM_UNHEALTHY', + restoreResult === 'push_failed' + ? 'Checkpoint restore left the workspace in an unknown state' + : 'Checkpoint fetch failed; refusing to run against an empty workspace', + ); + } + restored = restoreResult === 'restored'; + } + + const runningOk = await writeRuntimeSessionRecord( + { ...launchedRecord, state: 'RUNNING', last_seen_at: Date.now() }, + lockToken, + ); + if (!runningOk) { + throw new SandboxBackendError( + 'MICROVM_FENCED', + `Lost session lock for ${runtimeSessionId} after restore`, + ); + } + return { vm, restored }; + } catch (error) { + const terminated = await this.terminate(client, vm.microvmId, 'error'); + if (terminated) { + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + } + throw error; + } + } + + /** Polls the runner's health endpoint until it responds, bounded by the + * launch timeout — a freshly-RUNNING VM's app may still be booting. */ + private async waitForRunnerReady( + client: LambdaMicrovmClient, + microvmId: string, + base: string, + ctx: SandboxExecuteContext, + ): Promise { + const deadline = Date.now() + this.config.launchTimeoutMs; + const refreshSkewMs = this.config.healthTimeoutMs + this.pollIntervalMs; + let token: MicrovmAuthToken | undefined; + let lastError: unknown; + while (Date.now() < deadline) { + ctx.signal.throwIfAborted(); + /* Readiness can legitimately outlive one proxy token under a custom + * launch timeout. Refresh just before expiry instead of repeatedly + * probing with a credential AWS will reject. A newly minted token is + * still used once even when a test/fake gives it an unusually short + * lifetime, avoiding a mint-only loop. */ + if (token === undefined || token.expiresAtMs <= Date.now() + refreshSkewMs) { + token = await this.mintAuthToken(client, microvmId, ctx.signal); + } + try { + await this.assertHealthy(base, token.token, ctx); + return; + } catch (error) { + ctx.signal.throwIfAborted(); + lastError = error; + if (Date.now() + this.pollIntervalMs >= deadline) break; + await sleep(this.pollIntervalMs, ctx.signal); + } + } + throw lastError instanceof SandboxBackendError + ? lastError + : new SandboxBackendError('MICROVM_UNHEALTHY', 'Runner did not become ready before restore', lastError); + } + + /** + * Pulls a checkpoint from the still-warm VM while the exec lock is held and + * stores it, returning the record to persist (with the checkpoint pointer) + * or the liveness-only update if checkpoints are off/failed. Never throws — + * a missed post-exec checkpoint does not rewrite the prior durable pointer. + * The current warm VM remains authoritative until the next successful + * checkpoint; a later replacement restores the last committed snapshot. + */ + private async checkpointUnderLock( + client: LambdaMicrovmClient, + record: RuntimeSessionRecord, + runtimeSessionId: string, + now: number, + lockToken: string, + signal?: AbortSignal, + ): Promise { + const base: RuntimeSessionRecord = { ...record, state: 'RUNNING', last_seen_at: now }; + if (!this.checkpointStore || !this.checkpointsActive() || !record.microvm_id || !record.endpoint) { + return base; + } + const result = await checkpointSession({ + mintToken: (microvmId) => this.mintAuthToken(client, microvmId, signal), + store: this.checkpointStore, + runtimeSessionId, + config: this.config.checkpoint, + normalizeEndpoint: normalizeMicrovmEndpoint, + lockToken, + signal, + }); + /* checkpointSession wrote the pointer under our lock on success; re-read so + * we keep it, but re-apply `last_seen_at: now` — that record was built from + * the pre-execute snapshot, so without this the liveness timestamp never + * advances on checkpointed executes and an actively-used session would look + * idle and relaunch needlessly. */ + if (result === 'stored') { + const persisted = await readRuntimeSessionRecord(runtimeSessionId); + return persisted ? { ...persisted, last_seen_at: now } : base; + } + return base; + } + + /** Mints a proxy auth token under the shared per-second token budget, so + * concurrent warm-session executes queue instead of bursting past AWS's + * CreateMicrovmAuthToken TPS limit (mirrors launch's `run` budget). Maps + * control-plane failures to SandboxBackendError so they never escape raw: + * `throttled` poisons the bucket for backoff, and `not_found` (the VM was + * evicted/terminated) surfaces as MICROVM_UNHEALTHY so the caller tears down + * the stale record and relaunches instead of retrying a dead VM. */ + private async mintAuthToken( + client: LambdaMicrovmClient, + microvmId: string, + signal?: AbortSignal, + ): Promise { + try { + await acquireOpBudget('token', { + limitPerSecond: this.config.tokenTps, + budgetMs: this.config.launchTimeoutMs, + signal, + }); + } catch (error) { + if (error instanceof MicrovmOpThrottledError) { + microvmThrottleEvents.inc({ op: 'token' }); + throw new SandboxBackendError('MICROVM_LAUNCH_THROTTLED', error.message, error); + } + throw error; + } + try { + return await client.createMicrovmAuthToken({ + microvmId, + port: this.config.port, + ttlSeconds: this.config.authTokenTtlSeconds, + }, signal); + } catch (error) { + if (error instanceof LambdaMicrovmApiError && error.kind === 'throttled') { + await poisonOpBucket('token'); + microvmThrottleEvents.inc({ op: 'token' }); + throw new SandboxBackendError('MICROVM_LAUNCH_THROTTLED', error.message, error); + } + if (error instanceof LambdaMicrovmApiError && error.kind === 'not_found') { + throw new SandboxBackendError('MICROVM_UNHEALTHY', error.message, error); + } + if (error instanceof LambdaMicrovmApiError) { + throw new SandboxBackendError('MICROVM_LAUNCH_FAILED', error.message, error); + } + throw error; + } + } + + /** + * Ensures the VM holds every by-reference input this request declares, + * BEFORE the execute. The guest cannot reach the file server, so the control + * plane fetches the bytes and pushes them into the runner's input cache; the + * runner's normal priming then resolves refs from that cache. + * + * Dedupe is a probe, not bookkeeping: the VM reports what it is missing, so + * a lost session record (recycle, failover, flush) can never cause a needed + * input to be skipped — and a re-push can never revert a sandbox edit, + * because this path does not write the workspace at all. + * + * Applies to session AND stateless executions alike: the cache is keyed by + * (storage session, object), not by session. + */ + private async deliverInputs( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + onUnhealthy?: () => Promise, + ): Promise { + const refs = req.inputDelivery ?? sessionFileRefs(req.body.files); + if (refs.length === 0) return; + if (refs.length > SESSION_INPUTS_MAX_COUNT) { + throw new SessionFilesError( + 'SESSION_INPUT_TOO_LARGE', + `Session delivery of ${refs.length} objects exceeds the ${SESSION_INPUTS_MAX_COUNT} limit`, + ); + } + const endpointBase = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + const mintToken = () => this.mintAuthToken(client, vm.microvmId, ctx.signal); + const deliveryConfig = { + ...this.config.checkpoint, + /* A probe can be the request that auto-resumes a suspended VM. Give that + * operation the worker's full budget; the shared abort signal still + * enforces the actual remaining JOB_TIMEOUT. */ + timeoutMs: this.config.jobTimeoutMs, + }; + let missing: Array<{ cache_key: string }>; + try { + missing = await this.probeInputsWithRetry( + { mintToken, endpointBase, signal: ctx.signal }, + refs.map((ref) => ({ cache_key: ref.cache_key })), + deliveryConfig, + ); + } catch (error) { + return await this.rethrowInputDeliveryFailure( + error, + ctx.signal, + onUnhealthy, + 'Session input delivery failed', + ); + } + logger.info('Session input delivery', { + microvmId: vm.microvmId, + refs: refs.length, + missing: missing.length, + }); + if (missing.length === 0) return; + + /* Source fetch/authorization/size failures say nothing about VM health. + * Keep a warm session intact so a caller error cannot discard workspace + * changes newer than the last checkpoint. */ + const wanted = new Set(missing.map((ref) => ref.cache_key)); + const batch = await buildInputBatch( + refs.filter((ref) => wanted.has(ref.cache_key)), + { + timeoutMs: this.config.checkpoint.timeoutMs, + maxBytes: this.config.checkpoint.maxBytes, + signal: ctx.signal, + }, + ); + if (!batch) { + throw new SessionFilesError( + 'SESSION_INPUT_PREPARATION_FAILED', + 'Runner requested an empty session input batch', + ); + } + + try { + await pushInputs( + { mintToken, endpointBase, signal: ctx.signal }, + () => fs.createReadStream(batch.path), + deliveryConfig, + batch.size, + batch.expandedSize, + ); + } catch (error) { + return await this.rethrowInputDeliveryFailure( + error, + ctx.signal, + onUnhealthy, + 'Session input push failed', + ); + } finally { + await batch.cleanup().catch(() => {}); + } + + /* Pruning happens as part of the push. Re-probe the ENTIRE working set so + * an undersized runner cache cannot ACK the batch and then execute with an + * older, currently-required ref evicted. */ + let stillMissing: Array<{ cache_key: string }>; + try { + stillMissing = await this.probeInputsWithRetry( + { mintToken, endpointBase, signal: ctx.signal }, + refs.map((ref) => ({ cache_key: ref.cache_key })), + deliveryConfig, + ); + } catch (error) { + return await this.rethrowInputDeliveryFailure( + error, + ctx.signal, + onUnhealthy, + 'Session input verification failed', + ); + } + if (stillMissing.length > 0) { + throw new SessionFilesError( + 'SESSION_INPUT_TOO_LARGE', + `Runner input cache cannot hold the ${refs.length}-object working set`, + ); + } + } + + /** + * Input delivery is control-plane/cache work and has not started an NsJail + * execution. Recycle a warm VM only when the failure positively says that VM + * is gone or unreachable. Token throttles, other Lambda API failures, live + * runner 4xx/5xx responses, validation errors, and caller aborts do not make + * the VM unsafe to reuse. + */ + private inputDeliveryProvesVmUnhealthy(error: unknown, signal: AbortSignal): boolean { + if (signal.aborted) return false; + if (error instanceof SandboxBackendError) { + return error.code === 'MICROVM_UNHEALTHY'; + } + if (!axios.isAxiosError(error)) return false; + if (error.response == null) return true; + return error.response.status >= 502 && error.response.status <= 504; + } + + private async rethrowInputDeliveryFailure( + error: unknown, + signal: AbortSignal, + onUnhealthy: (() => Promise) | undefined, + message: string, + ): Promise { + if (this.inputDeliveryProvesVmUnhealthy(error, signal)) { + await onUnhealthy?.(); + } + /* Preserve mapped control-plane codes (especially throttling) so callers + * can apply their normal retry/backpressure policy. */ + if (error instanceof SandboxBackendError) { + throw error; + } + throw new SandboxBackendError('MICROVM_UNHEALTHY', message, error); + } + + private async probeInputsWithRetry( + args: { mintToken: () => Promise; endpointBase: string; signal?: AbortSignal }, + refs: Array<{ cache_key: string }>, + config: CheckpointConfig, + ): Promise> { + const deadline = Date.now() + config.timeoutMs; + for (;;) { + try { + return await probeInputs(args, refs, config); + } catch (error) { + const status = axios.isAxiosError(error) ? error.response?.status ?? 0 : 0; + /* A suspended endpoint can reset/refuse the connection while AWS + * auto-resumes it, before the proxy is ready to return a 502/503/504. + * Retry both transport-level failures and gateway transients under the + * caller's full bounded delivery budget. */ + const transient = + axios.isAxiosError(error) + && (error.response == null || status === 502 || status === 503 || status === 504); + if (!transient || args.signal?.aborted || Date.now() + this.pollIntervalMs >= deadline) { + throw error; + } + await sleep(this.pollIntervalMs, args.signal); + } + } + } + + private async proxyExecute( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + runtimeSessionId?: string, + readinessSatisfied = false, + ): Promise { + const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + const token = await this.mintAuthToken(client, vm.microvmId, ctx.signal); + /* Skip the preflight health check on a REUSED session VM: AWS may be + * auto-resuming it from a suspend, and that resume can exceed the short + * healthTimeoutMs (it scales with suspended-state size) — a slow-but-valid + * resume would be misread as MICROVM_UNHEALTHY and the VM torn down. The + * execute itself carries the resume under the full job budget, a + * genuinely-evicted VM already fails token minting with not_found, and a + * freshly-launched VM (reused=false) still gets the readiness probe. */ + if (!readinessSatisfied) { + /* A freshly-launched VM's runner may still be booting (RUNNING is a + * control-plane state). When checkpoints are active the launch path + * already polled readiness; when they are disabled/stateless this is the + * only gate — so poll until launchTimeoutMs rather than tearing the VM + * down on a single connection-refused/503 probe. */ + await this.waitForRunnerReady(client, vm.microvmId, base, ctx); + } + + /* Session mode is opted into per-request via this header (not a /run + * lifecycle hook): stock container images can't route Lambda's build + * hooks, so the runner reads its runtime session id straight off the + * proxied execute. Header-only — never the manifest-signed body. */ + const sessionHeader = runtimeSessionId + ? { [RUNTIME_SESSION_ID_HEADER]: runtimeSessionId } + : undefined; + + return withSpan('codeapi.sandbox.execute', { + 'http.request.method': 'POST', + 'url.path': `/${Jobs.execute}`, + 'codeapi.language': ctx.language, + 'codeapi.sandbox.backend': this.name, + }, async () => { + const response = await axios.post( + `${base}/api/v2/${Jobs.execute}`, + req.body, + { + headers: { + ...injectTraceHeaders(req.headers), + [token.headerName]: token.token, + ...microvmPortHeaders(this.config.port), + ...sessionHeader, + }, + signal: ctx.signal, + }, + ); + if (response.status !== 200) { + throw new Error('Error from sandbox'); + } + return response.data; + }, 'CLIENT'); + } + + private async launch( + client: LambdaMicrovmClient, + ctx: SandboxExecuteContext, + opts: LaunchOptions, + ): Promise { + try { + return await this.launchOnce(client, ctx, opts); + } catch (error) { + const retryable = + error instanceof SandboxBackendError && + error.code === 'MICROVM_LAUNCH_FAILED' && + error.transient && + !ctx.signal.aborted; + if (!retryable) { + throw error; + } + /* Fresh clientToken: RunMicrovm is idempotent per token, so reusing the + * original would hand back the same dead VM. */ + logger.warn( + `[${ctx.executionId}] MicroVM died during boot (${error.message}); retrying launch once`, + ); + microvmLaunches.inc({ outcome: 'retried' }); + return this.launchOnce(client, ctx, { ...opts, clientToken: `${opts.clientToken}-r1` }); + } + } + + private async launchOnce( + client: LambdaMicrovmClient, + ctx: SandboxExecuteContext, + opts: LaunchOptions, + ): Promise { + const endLaunchTimer = microvmLaunchDuration.startTimer(); + try { + await acquireOpBudget('run', { + limitPerSecond: this.config.launchTps, + budgetMs: this.config.launchTimeoutMs, + signal: ctx.signal, + }); + } catch (error) { + if (error instanceof MicrovmOpThrottledError) { + microvmThrottleEvents.inc({ op: 'run' }); + microvmLaunches.inc({ outcome: 'throttled' }); + throw new SandboxBackendError('MICROVM_LAUNCH_THROTTLED', error.message, error); + } + throw error; + } + + let vm: MicrovmDescription; + try { + vm = await client.runMicrovm({ + imageIdentifier: this.config.imageArn, + imageVersion: this.config.imageVersion, + executionRoleArn: this.config.executionRoleArn, + logGroup: this.config.logGroup, + ingressConnectorArns: this.config.ingressConnectorArns, + egressConnectorArns: this.config.egressConnectorArns, + maximumDurationSeconds: opts.maxDurationSeconds, + idlePolicy: opts.idlePolicy, + clientToken: opts.clientToken, + }, ctx.signal); + } catch (error) { + if (error instanceof LambdaMicrovmApiError && error.kind === 'throttled') { + await poisonOpBucket('run'); + microvmThrottleEvents.inc({ op: 'run' }); + microvmLaunches.inc({ outcome: 'throttled' }); + throw new SandboxBackendError('MICROVM_LAUNCH_THROTTLED', error.message, error); + } + microvmLaunches.inc({ outcome: 'failed' }); + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + error instanceof Error ? error.message : 'RunMicrovm failed', + error, + ); + } + + try { + const ready = await this.waitUntilRunning(client, vm, ctx); + microvmLaunches.inc({ outcome: 'ok' }); + endLaunchTimer(); + return ready; + } catch (error) { + microvmLaunches.inc({ outcome: 'failed' }); + const terminated = await this.terminate(client, vm.microvmId, 'error'); + if (!terminated) { + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + `MicroVM ${vm.microvmId} failed during boot and could not be terminated`, + error, + ); + } + /* waitUntilRunning throws SandboxBackendError for its own conditions, but + * the GetMicrovm poll it makes can throw a raw LambdaMicrovmApiError + * (throttle/transient control-plane error). Map it like runMicrovm so it + * surfaces as a public MICROVM_LAUNCH_* failure, not a generic 500. */ + if (error instanceof SandboxBackendError) throw error; + if (error instanceof LambdaMicrovmApiError && error.kind === 'throttled') { + await poisonOpBucket('run'); + microvmThrottleEvents.inc({ op: 'run' }); + throw new SandboxBackendError('MICROVM_LAUNCH_THROTTLED', error.message, error); + } + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + error instanceof Error ? error.message : 'MicroVM poll failed', + error, + ); + } + } + + private async waitUntilRunning( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + ctx: SandboxExecuteContext, + ): Promise { + const deadline = Date.now() + this.config.launchTimeoutMs; + let current = vm; + while (current.state !== 'RUNNING' || !current.endpoint) { + if (ctx.signal.aborted) { + throw new SandboxBackendError('MICROVM_LAUNCH_FAILED', 'Execution aborted while MicroVM was launching'); + } + if (current.state === 'TERMINATED' || current.state === 'TERMINATING') { + /* A boot-time death is a fast, provider-side transient (observed in + * the field a few times a day) — mark it retryable so `launch` can + * try once more instead of surfacing a 503 to the caller. */ + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + `MicroVM ${current.microvmId} entered ${current.state} before becoming ready`, + undefined, + true, + ); + } + if (Date.now() + this.pollIntervalMs > deadline) { + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + `MicroVM ${current.microvmId} did not reach RUNNING within ${this.config.launchTimeoutMs}ms`, + ); + } + await sleep(this.pollIntervalMs, ctx.signal); + current = await client.getMicrovm(current.microvmId, ctx.signal); + } + return current; + } + + private async assertHealthy(base: string, token: string, ctx: SandboxExecuteContext): Promise { + try { + const response = await axios.get(`${base}/api/v2/health`, { + headers: { 'X-aws-proxy-auth': token, ...microvmPortHeaders(this.config.port) }, + timeout: this.config.healthTimeoutMs, + signal: ctx.signal, + }); + if (response.status !== 200) { + throw new Error(`health returned ${response.status}`); + } + } catch (error) { + throw new SandboxBackendError( + 'MICROVM_UNHEALTHY', + `MicroVM health check failed: ${error instanceof Error ? error.message : 'unknown error'}`, + error, + ); + } + } + + private async terminate( + client: LambdaMicrovmClient, + microvmId: string, + reason: string, + ): Promise { + try { + await client.terminateMicrovm( + microvmId, + AbortSignal.timeout(this.config.launchTimeoutMs), + ); + microvmTerminations.inc({ reason }); + return true; + } catch (error) { + if (error instanceof LambdaMicrovmApiError && error.kind === 'not_found') { + /* The desired terminal state already holds. Treat provider not-found as + * successful cleanup so a stale registry record cannot wedge relaunch. */ + microvmTerminations.inc({ reason }); + return true; + } + logger.error('Failed to terminate MicroVM', { microvmId, reason, error }); + return false; + } + } +} diff --git a/service/src/sandbox-backend/types.ts b/service/src/sandbox-backend/types.ts new file mode 100644 index 0000000..424bc1a --- /dev/null +++ b/service/src/sandbox-backend/types.ts @@ -0,0 +1,101 @@ +import type * as t from '../types'; +import { getAxiosErrorDetails } from '../utils'; + +/** + * Fully built sandbox execute request: `body` already carries the egress + * grant and signed execution manifest from `buildSandboxExecuteRequest`. + * Backends MUST NOT mutate `body` — the manifest binds its sha256. + */ +export interface SandboxTransportRequest { + body: t.PayloadBody; + headers: Record; + /** + * Control-plane-only fetch refs. These retain authorized storage IDs while + * `body.files` may contain per-grant sandbox handles. Backends must never + * serialize this field to the sandbox. + */ + inputDelivery?: SandboxInputDeliveryRef[]; +} + +export interface SandboxInputDeliveryRef { + id: string; + storage_session_id: string; + name: string; + cache_key: string; +} + +export interface SandboxExecuteContext { + executionId: string; + language: string; + isSynthetic: boolean; + /** Worker-owned JOB_TIMEOUT abort signal. */ + signal: AbortSignal; + /** Absolute deadline anchored to BullMQ enqueue time when available. Unlike + * a backend-local timer, this includes queue wait, egress-grant creation, and + * request setup. Direct backend callers may omit it and receive a + * backend-entry fallback. */ + deadlineAtMs?: number; + tenantId?: string; + canonicalUserId?: string; + /** Absent ⇒ stateless execution (no runtime session affinity). */ + runtimeSessionId?: string; + runtimeSessionMode: 'stateless' | 'affinity' | 'strict'; + /** + * Worker-owned result transformation that must succeed before a stateful + * backend makes workspace mutations reusable or durable. Stateless backends + * intentionally leave this to the worker's normal post-execute path. + */ + sessionResultFinalizer?: (result: SandboxRawResponse) => Promise; +} + +/** Raw sandbox response, pre-gateway-restore. */ +export type SandboxRawResponse = t.ExecuteResponse & { + session_id: string; + files?: t.FileRefs; + run?: t.ExecuteResponse['run']; +}; + +export interface SandboxBackend { + readonly name: 'http' | 'lambda-microvm'; + execute(req: SandboxTransportRequest, ctx: SandboxExecuteContext): Promise; + shutdown?(): Promise; +} + +export type SandboxBackendErrorCode = + | 'RUNTIME_SESSION_BUSY' + | 'MICROVM_LAUNCH_FAILED' + | 'MICROVM_LAUNCH_THROTTLED' + | 'MICROVM_UNHEALTHY' + | 'MICROVM_FENCED' + | 'MICROVM_DEADLINE_EXCEEDED'; + +/** Lambda-only failure modes; the worker prefixes messages with the code so + * the router can map them (e.g. RUNTIME_SESSION_BUSY -> 409). Axios errors + * from the sandbox POST itself are rethrown raw by every backend. */ +export class SandboxBackendError extends Error { + /** + * The originating failure, ALWAYS stored sanitized. Backend causes are + * routinely axios errors whose `config` carries the minted MicroVM auth + * header, the internal service token, and (on a push) the request body — + * i.e. the archive bytes. Callers log wrapper errors wholesale, so + * sanitizing at construction is the only place that covers every path. + */ + public readonly cause?: unknown; + + /** + * @param transient - Marks a failure the backend may safely retry once with + * fresh identifiers (e.g. a MicroVM that reached a terminal state during + * boot). Throttles, aborts, and deadline timeouts stay non-transient: a + * throttle retry worsens the pressure and a timeout retry doubles the wait. + */ + constructor( + public readonly code: SandboxBackendErrorCode, + message: string, + cause?: unknown, + public readonly transient: boolean = false, + ) { + super(message); + this.name = 'SandboxBackendError'; + this.cause = cause === undefined ? undefined : getAxiosErrorDetails(cause); + } +} diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 15224ad..c4937b0 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -3,12 +3,34 @@ import { env } from './config'; import { validateApiHardenedConfig, validateEgressGatewayHardenedConfig, + validateSandboxBackendPolicy, validateWorkerHardenedConfig, } from './secure-startup'; const savedEnv = { ...process.env }; const saved = { hardened: env.HARDENED_SANDBOX_MODE, + sandboxBackend: env.SANDBOX_BACKEND, + ptcMode: env.PTC_MODE, + runtimeSessionMode: env.RUNTIME_SESSION_MODE, + lambdaImageArn: env.LAMBDA_MICROVM_IMAGE_ARN, + lambdaImageVersion: env.LAMBDA_MICROVM_IMAGE_VERSION, + lambdaEgressConnectors: env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS, + lambdaPort: env.LAMBDA_MICROVM_PORT, + lambdaMaxDuration: env.LAMBDA_MICROVM_MAX_DURATION_SECONDS, + lambdaIdleSeconds: env.LAMBDA_MICROVM_IDLE_SECONDS, + lambdaSuspendSeconds: env.LAMBDA_MICROVM_SUSPEND_SECONDS, + lambdaTokenTtl: env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS, + lambdaLaunchTimeout: env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, + lambdaHealthTimeout: env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS, + lambdaLaunchTps: env.LAMBDA_MICROVM_LAUNCH_TPS, + lambdaTokenTps: env.LAMBDA_MICROVM_TOKEN_TPS, + lambdaAllowShell: env.LAMBDA_MICROVM_ALLOW_SHELL, + jobTimeout: env.JOB_TIMEOUT, + lockWait: env.RUNTIME_SESSION_LOCK_WAIT_MS, + checkpoints: env.SESSION_CHECKPOINTS, + checkpointMaxBytes: env.CHECKPOINT_MAX_BYTES, + checkpointTimeout: env.CHECKPOINT_TIMEOUT_MS, gatewayUrl: env.EGRESS_GATEWAY_URL, grantSecret: env.EGRESS_GRANT_SECRET, privateKey: env.EXECUTION_MANIFEST_PRIVATE_KEY, @@ -24,6 +46,27 @@ function restore(): void { } Object.assign(process.env, savedEnv); env.HARDENED_SANDBOX_MODE = saved.hardened; + env.SANDBOX_BACKEND = saved.sandboxBackend; + env.PTC_MODE = saved.ptcMode; + env.RUNTIME_SESSION_MODE = saved.runtimeSessionMode; + env.LAMBDA_MICROVM_IMAGE_ARN = saved.lambdaImageArn; + env.LAMBDA_MICROVM_IMAGE_VERSION = saved.lambdaImageVersion; + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = saved.lambdaEgressConnectors; + env.LAMBDA_MICROVM_PORT = saved.lambdaPort; + env.LAMBDA_MICROVM_MAX_DURATION_SECONDS = saved.lambdaMaxDuration; + env.LAMBDA_MICROVM_IDLE_SECONDS = saved.lambdaIdleSeconds; + env.LAMBDA_MICROVM_SUSPEND_SECONDS = saved.lambdaSuspendSeconds; + env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS = saved.lambdaTokenTtl; + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS = saved.lambdaLaunchTimeout; + env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS = saved.lambdaHealthTimeout; + env.LAMBDA_MICROVM_LAUNCH_TPS = saved.lambdaLaunchTps; + env.LAMBDA_MICROVM_TOKEN_TPS = saved.lambdaTokenTps; + env.LAMBDA_MICROVM_ALLOW_SHELL = saved.lambdaAllowShell; + env.JOB_TIMEOUT = saved.jobTimeout; + env.RUNTIME_SESSION_LOCK_WAIT_MS = saved.lockWait; + env.SESSION_CHECKPOINTS = saved.checkpoints; + env.CHECKPOINT_MAX_BYTES = saved.checkpointMaxBytes; + env.CHECKPOINT_TIMEOUT_MS = saved.checkpointTimeout; env.EGRESS_GATEWAY_URL = saved.gatewayUrl; env.EGRESS_GRANT_SECRET = saved.grantSecret; env.EXECUTION_MANIFEST_PRIVATE_KEY = saved.privateKey; @@ -120,3 +163,195 @@ describe('hardened CodeAPI startup config', () => { expect(() => validateEgressGatewayHardenedConfig()).toThrow('CODEAPI_EGRESS_LEDGER_REQUIRED'); }); }); + +describe('sandbox backend policy', () => { + function configureValidLambda(): void { + env.SANDBOX_BACKEND = 'lambda-microvm'; + env.HARDENED_SANDBOX_MODE = false; + env.PTC_MODE = 'replay'; + env.RUNTIME_SESSION_MODE = 'stateless'; + env.LAMBDA_MICROVM_IMAGE_ARN = 'arn:aws:lambda:us-east-2:1:microvm-image:codeapi'; + env.LAMBDA_MICROVM_IMAGE_VERSION = '3'; + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = undefined; + env.LAMBDA_MICROVM_PORT = 8080; + env.LAMBDA_MICROVM_MAX_DURATION_SECONDS = 28_800; + env.LAMBDA_MICROVM_IDLE_SECONDS = 1_800; + env.LAMBDA_MICROVM_SUSPEND_SECONDS = 1_800; + env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS = 300; + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS = 60_000; + env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS = 5_000; + env.LAMBDA_MICROVM_LAUNCH_TPS = 4; + env.LAMBDA_MICROVM_TOKEN_TPS = 8; + env.LAMBDA_MICROVM_ALLOW_SHELL = false; + env.JOB_TIMEOUT = 300_000; + env.RUNTIME_SESSION_LOCK_WAIT_MS = 15_000; + env.SESSION_CHECKPOINTS = true; + env.CHECKPOINT_MAX_BYTES = 512 * 1024 * 1024; + env.CHECKPOINT_TIMEOUT_MS = 60_000; + /* Object storage for the (default-on) session checkpoints. */ + process.env.MINIO_ENDPOINT = 'minio'; + process.env.MINIO_ACCESS_KEY = 'access'; + process.env.MINIO_SECRET_KEY = 'secret'; + process.env.CODEAPI_CHECKPOINT_BUCKET = 'codeapi-checkpoints'; + } + + test('accepts the default http backend', () => { + env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'stateless'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('accepts a fully configured stateless lambda backend', () => { + configureValidLambda(); + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('strict runtime sessions require the lambda backend', () => { + env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'strict'; + expect(() => validateSandboxBackendPolicy()).toThrow('requires the lambda-microvm backend'); + }); + + test('rejects blocking PTC on the lambda backend', () => { + configureValidLambda(); + env.PTC_MODE = 'blocking'; + expect(() => validateSandboxBackendPolicy()).toThrow('PTC replay is the only supported PTC mode'); + }); + + test('requires the image ARN', () => { + configureValidLambda(); + env.LAMBDA_MICROVM_IMAGE_ARN = ''; + expect(() => validateSandboxBackendPolicy()).toThrow('LAMBDA_MICROVM_IMAGE_ARN is required'); + }); + + test('accepts affinity and strict session modes on the lambda backend', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + env.RUNTIME_SESSION_MODE = 'strict'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('accepts the production checkpoint budget and rejects an impossible one', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + /* Production defaults reserve 60s for token throttling, 60s for each + * archive transfer, and 5s for each metadata operation: 190s total. */ + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + + env.JOB_TIMEOUT = 190_000; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'JOB_TIMEOUT must exceed the full session checkpoint reserve (190000ms)', + ); + }); + + test('stateful session modes require a pinned image version', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + env.LAMBDA_MICROVM_IMAGE_VERSION = undefined; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'LAMBDA_MICROVM_IMAGE_VERSION must be pinned', + ); + }); + + test('rejects session checkpoints without object storage configured', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + delete process.env.MINIO_ENDPOINT; + expect(() => validateSandboxBackendPolicy()).toThrow('object storage is not configured'); + /* stateless never touches the store, so it stays valid without MinIO */ + env.RUNTIME_SESSION_MODE = 'stateless'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('accepts IAM-role checkpoint credentials and rejects a partial static pair', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + delete process.env.MINIO_ACCESS_KEY; + delete process.env.MINIO_SECRET_KEY; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + + process.env.MINIO_ACCESS_KEY = 'incomplete'; + expect(() => validateSandboxBackendPolicy()).toThrow('MINIO_SECRET_KEY'); + }); + + test('hardened mode requires an egress connector', () => { + configureValidLambda(); + env.HARDENED_SANDBOX_MODE = true; + expect(() => validateSandboxBackendPolicy()).toThrow('LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS is required'); + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = ['arn:aws:lambda:us-east-2:1:network-connector:vpc-egress']; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('hardened mode rejects public INTERNET_EGRESS, including mixed connector lists', () => { + configureValidLambda(); + env.HARDENED_SANDBOX_MODE = true; + const publicEgress = + 'arn:aws:lambda:us-east-2:aws:network-connector:aws-network-connector:INTERNET_EGRESS'; + const vpcEgress = 'arn:aws:lambda:us-east-2:1:network-connector:vpc-egress'; + + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = [publicEgress]; + expect(() => validateSandboxBackendPolicy()).toThrow('INTERNET_EGRESS'); + + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = [vpcEgress, publicEgress]; + expect(() => validateSandboxBackendPolicy()).toThrow('INTERNET_EGRESS'); + + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = [vpcEgress]; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('rejects invalid Lambda numeric config and preserves suspend=0', () => { + configureValidLambda(); + env.LAMBDA_MICROVM_IDLE_SECONDS = 59; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'LAMBDA_MICROVM_IDLE_SECONDS must be a whole number between 60 and 28800', + ); + + configureValidLambda(); + env.LAMBDA_MICROVM_MAX_DURATION_SECONDS = 28_801; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'LAMBDA_MICROVM_MAX_DURATION_SECONDS must be a whole number between 1 and 28800', + ); + + configureValidLambda(); + env.LAMBDA_MICROVM_SUSPEND_SECONDS = 0; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('rejects unbounded or invalid session deadline and checkpoint controls', () => { + configureValidLambda(); + env.RUNTIME_SESSION_LOCK_WAIT_MS = Number.POSITIVE_INFINITY; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'CODEAPI_RUNTIME_SESSION_LOCK_WAIT_MS', + ); + + configureValidLambda(); + env.CHECKPOINT_MAX_BYTES = Number.POSITIVE_INFINITY; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'CODEAPI_CHECKPOINT_MAX_BYTES', + ); + + configureValidLambda(); + env.CHECKPOINT_TIMEOUT_MS = 0; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'CODEAPI_CHECKPOINT_TIMEOUT_MS', + ); + + configureValidLambda(); + env.JOB_TIMEOUT = Number.NaN; + expect(() => validateSandboxBackendPolicy()).toThrow('JOB_TIMEOUT'); + }); + + test('caps ingress token TTL and blocks shell ingress in hardened mode', () => { + configureValidLambda(); + env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS = 901; + expect(() => validateSandboxBackendPolicy()).toThrow('must be a whole number between 1 and 900'); + + configureValidLambda(); + env.LAMBDA_MICROVM_ALLOW_SHELL = true; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + env.HARDENED_SANDBOX_MODE = true; + env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS = ['arn:aws:lambda:us-east-2:1:network-connector:vpc-egress']; + expect(() => validateSandboxBackendPolicy()).toThrow('LAMBDA_MICROVM_ALLOW_SHELL'); + }); +}); diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 1d84e90..cf0a78d 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -1,4 +1,9 @@ -import { env } from './config'; +import { + checkpointPipelineBudgetMs, + env, + lambdaMicrovmNumericConfigError, +} from './config'; +import logger from './logger'; import { INTERNAL_SERVICE_TOKEN_ENV } from './internal-service-auth'; export class SecureStartupConfigError extends Error { @@ -31,6 +36,17 @@ function rejectValue(name: string, value: string | undefined): void { } } +function requireSafeWholeNumber(name: string, value: number, min: number): void { + if (!Number.isSafeInteger(value) || value < min) { + throw new SecureStartupConfigError( + `${name} must be a whole number of at least ${min}`, + ); + } +} + +const AWS_MANAGED_INTERNET_EGRESS_SUFFIX = + ':aws:network-connector:aws-network-connector:INTERNET_EGRESS'; + export function validateApiHardenedConfig(): void { if (!env.HARDENED_SANDBOX_MODE) return; rejectValue('CODEAPI_EGRESS_GRANT_SECRET', process.env.CODEAPI_EGRESS_GRANT_SECRET); @@ -48,6 +64,113 @@ export function validateWorkerHardenedConfig(): void { requireValue('CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY', env.EXECUTION_MANIFEST_PRIVATE_KEY); } +/** + * Backend-selection policy. Unlike the hardened-mode validators, this runs + * unconditionally: a misconfigured backend must never half-start. + */ +export function validateSandboxBackendPolicy(): void { + if (env.RUNTIME_SESSION_MODE === 'strict' && env.SANDBOX_BACKEND === 'http') { + throw new SecureStartupConfigError( + 'CODEAPI_RUNTIME_SESSION_MODE=strict requires the lambda-microvm backend', + ); + } + /* affinity+http is the documented graceful fallback (runs stateless), but a + * silent no-op hides a likely misconfiguration — surface it loudly. */ + if (env.RUNTIME_SESSION_MODE === 'affinity' && env.SANDBOX_BACKEND === 'http') { + logger.warn( + 'CODEAPI_RUNTIME_SESSION_MODE=affinity has no effect with the http backend: ' + + 'requests run stateless (no session reuse or checkpoint/restore). ' + + 'Use CODEAPI_SANDBOX_BACKEND=lambda-microvm for stateful sessions.', + ); + } + if (env.SANDBOX_BACKEND !== 'lambda-microvm') return; + + const numericConfigError = lambdaMicrovmNumericConfigError(env); + if (numericConfigError !== undefined) { + throw new SecureStartupConfigError(numericConfigError); + } + requireSafeWholeNumber('JOB_TIMEOUT', env.JOB_TIMEOUT, 1); + requireSafeWholeNumber( + 'CODEAPI_RUNTIME_SESSION_LOCK_WAIT_MS', + env.RUNTIME_SESSION_LOCK_WAIT_MS, + 0, + ); + requireSafeWholeNumber('CODEAPI_CHECKPOINT_MAX_BYTES', env.CHECKPOINT_MAX_BYTES, 1); + requireSafeWholeNumber('CODEAPI_CHECKPOINT_TIMEOUT_MS', env.CHECKPOINT_TIMEOUT_MS, 1); + if (env.PTC_MODE === 'blocking') { + throw new SecureStartupConfigError( + 'PTC replay is the only supported PTC mode for the lambda-microvm backend (unset PTC_MODE=blocking)', + ); + } + if (env.LAMBDA_MICROVM_IMAGE_ARN.trim().length === 0) { + throw new SecureStartupConfigError('LAMBDA_MICROVM_IMAGE_ARN is required for the lambda-microvm backend'); + } + if ( + env.RUNTIME_SESSION_MODE !== 'stateless' + && !nonEmpty(env.LAMBDA_MICROVM_IMAGE_VERSION) + ) { + throw new SecureStartupConfigError( + 'LAMBDA_MICROVM_IMAGE_VERSION must be pinned for stateful runtime sessions', + ); + } + if (env.HARDENED_SANDBOX_MODE && (env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS?.length ?? 0) === 0) { + throw new SecureStartupConfigError( + 'LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS is required in CODEAPI_HARDENED_SANDBOX_MODE (MicroVMs default to public egress)', + ); + } + if ( + env.HARDENED_SANDBOX_MODE + && (env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS?.some( + (arn) => arn.trim().endsWith(AWS_MANAGED_INTERNET_EGRESS_SUFFIX), + ) ?? false) + ) { + throw new SecureStartupConfigError( + 'AWS-managed INTERNET_EGRESS is forbidden in CODEAPI_HARDENED_SANDBOX_MODE; use only a VPC egress connector restricted to the gateway', + ); + } + if (env.LAMBDA_MICROVM_ALLOW_SHELL && (env.HARDENED_SANDBOX_MODE || process.env.NODE_ENV === 'production')) { + throw new SecureStartupConfigError( + 'LAMBDA_MICROVM_ALLOW_SHELL must not be enabled in production or hardened mode', + ); + } + /* Checkpoints without object storage configured: MinioCheckpointStore silently + * falls back to localhost:9000/test-bucket, so warm reuse works + * but every checkpoint + restore fails against the dummy store and workspace + * state is lost on the first relaunch. Fail fast instead (mirrors the factory + * gate that constructs the store). Credentials may be a complete static + * MINIO pair or the workload's IAM role/web-identity provider. */ + if (env.SESSION_CHECKPOINTS && env.RUNTIME_SESSION_MODE !== 'stateless') { + const checkpointBudgetMs = checkpointPipelineBudgetMs( + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, + env.CHECKPOINT_TIMEOUT_MS, + ); + if (env.JOB_TIMEOUT <= checkpointBudgetMs) { + throw new SecureStartupConfigError( + `JOB_TIMEOUT must exceed the full session checkpoint reserve (${checkpointBudgetMs}ms): ` + + 'LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + two CODEAPI_CHECKPOINT_TIMEOUT_MS ' + + '+ two capped metadata operations', + ); + } + const missing = ['MINIO_ENDPOINT'].filter( + (name) => !nonEmpty(process.env[name]), + ); + const hasAccessKey = nonEmpty(process.env.MINIO_ACCESS_KEY); + const hasSecretKey = nonEmpty(process.env.MINIO_SECRET_KEY); + if (hasAccessKey !== hasSecretKey) { + missing.push(hasAccessKey ? 'MINIO_SECRET_KEY' : 'MINIO_ACCESS_KEY'); + } + if (!nonEmpty(process.env.CODEAPI_CHECKPOINT_BUCKET) && !nonEmpty(process.env.MINIO_BUCKET)) { + missing.push('CODEAPI_CHECKPOINT_BUCKET (or MINIO_BUCKET)'); + } + if (missing.length > 0) { + throw new SecureStartupConfigError( + 'Session checkpoints are enabled but object storage is not configured: ' + + `${missing.join(', ')}. Set them or disable CODEAPI_SESSION_CHECKPOINTS.`, + ); + } + } +} + export function validateEgressGatewayHardenedConfig(): void { if (!env.HARDENED_SANDBOX_MODE) return; rejectValue('CODEAPI_SYNTHETIC_ACCESS_TOKEN', process.env.CODEAPI_SYNTHETIC_ACCESS_TOKEN); diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 87bce41..51bccc7 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -15,6 +15,7 @@ import { internalServiceHeaders } from '../internal-service-auth'; import { resolveOutputBucketSessionKey, SessionKeyResolutionError } from '../session-key'; import { getCredentialId, getPrincipalOrReject } from '../auth/principal'; import { getExecutionIdentity } from '../execution-identity'; +import { PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION } from '../runtime-session/job-policy'; import { jobsSubmitted, ptcReplayContinuations, @@ -401,6 +402,7 @@ async function runReplayIteration( executionId: state.execution_id, tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, + runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, @@ -1368,6 +1370,7 @@ async function handleBlocking( executionId: execution_id, tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, + runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, diff --git a/service/src/service/router.ts b/service/src/service/router.ts index 9d9fd37..931974e 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -18,6 +18,7 @@ import { summarizeRequestedFiles } from '../execution-log'; import { getCredentialId, getPrincipalOrReject } from '../auth/principal'; import { isSyntheticPrincipalSource } from '../auth/synthetic'; import { getExecutionIdentity } from '../execution-identity'; +import { resolveRuntimeSessionIdForExecRequest, RuntimeSessionHintError } from '../runtime-session/id'; import { jobsSubmitted } from '../metrics'; import { captureTraceCarrier, withSpan } from '../telemetry'; import { Jobs, Languages } from '../enum'; @@ -133,6 +134,22 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) return res.status(400).json({ error: `Unsupported language: ${rawLang}` }); } + let runtimeSessionId: string | undefined; + try { + runtimeSessionId = resolveRuntimeSessionIdForExecRequest({ + mode: env.RUNTIME_SESSION_MODE, + storageNamespace: identity.storageNamespace, + canonicalUserId: identity.canonicalUserId, + runtimeSessionHint: body.runtime_session_hint, + isSynthetic: isSyntheticRequest, + }); + } catch (error) { + if (error instanceof RuntimeSessionHintError) { + return res.status(error.status).json({ error: error.message }); + } + throw error; + } + let authorizedFiles: t.RequestFile[]; try { authorizedFiles = await authorizeRequestedFiles({ @@ -219,6 +236,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) executionId: execution_id, tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, + ...(runtimeSessionId != null ? { runtimeSessionId } : {}), executionManifestClaims: sandboxSecurity.executionManifestClaims, egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, diff --git a/service/src/types/service.ts b/service/src/types/service.ts index e9e049f..d68af4f 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -128,6 +128,13 @@ export interface RequestBody { args?: string[]; user_id?: string; files?: RequestFile[]; + /** + * Optional stable identity hint (e.g. conversation id) for stateful + * runtime sessions. Never a security boundary: the server derives the + * final runtime session id as hash(tenant, user, hint). Ignored when + * CODEAPI_RUNTIME_SESSION_MODE is `stateless`. + */ + runtime_session_hint?: string; } export type CreatePayload = { req: AuthenticatedRequest, session_id: string; isPyPlot?: boolean }; @@ -150,6 +157,17 @@ export interface FileObject { } export type PayloadFile = { name: string; content: string }; +export type PayloadFileRef = { + id: string; + storage_session_id: string; + name: string; + /** + * Stable opaque identity for runner-local input caching. The worker derives + * it from the authorized raw ref before hardened egress replaces the raw + * identifiers with per-grant handles. + */ + input_cache_key?: string; +}; export interface PayloadBody { language: string; @@ -163,7 +181,7 @@ export interface PayloadBody { * are sessionKey-derivation inputs at the service entry, never * consumed by the sandbox, so they're intentionally not on this * shape. */ - files: Array; + files: Array; /** Top-level execution session id (passed to sandbox to seed Job.uuid). */ session_id?: string; /** Output storage session id/handle used for generated file uploads. */ @@ -216,6 +234,9 @@ export interface LanguageConfig { runtime?: string; } +/** Explicit reasons a queue job may remain stateless in a stateful mode. */ +export type RuntimeSessionExemption = 'programmatic'; + export type JobData = { code: string; userId: string; @@ -227,6 +248,16 @@ export type JobData = { executionId?: string; tenantId?: string; canonicalUserId?: string; + /** + * Server-derived runtime session identity. Absence is stateless unless + * strict mode requires it; explicit exemptions document intentional gaps. + */ + runtimeSessionId?: string; + /** + * Intentional stateless exception. Programmatic/replay execution externalizes + * continuation state and does not yet bind PTC rounds to a warm session. + */ + runtimeSessionExemption?: RuntimeSessionExemption; executionManifestClaims?: ExecutionManifestClaims; /** Raw grant claims retained only for service-worker dispatch so grant * expiry is anchored to sandbox start, not BullMQ enqueue time. */ diff --git a/service/src/utils.test.ts b/service/src/utils.test.ts index 34be93f..89dcd32 100644 --- a/service/src/utils.test.ts +++ b/service/src/utils.test.ts @@ -129,6 +129,86 @@ describe('sandbox error formatting', () => { }); }); + test('maps a busy runtime session (strict mode) to 409', () => { + const err = new Error('RUNTIME_SESSION_BUSY: Runtime session rt_abc is busy'); + expect(publicExecutionFailure(err)).toEqual({ + status: 409, + body: { error: 'runtime_session_busy', message: 'Runtime session is busy' }, + }); + }); + + test('maps MicroVM launch failures to 503', () => { + const err = new Error('MICROVM_LAUNCH_FAILED: MicroVM did not reach RUNNING within 60000ms'); + expect(publicExecutionFailure(err)).toEqual({ + status: 503, + body: { error: 'microvm_launch_failed', message: 'Sandbox launch failed' }, + }); + }); + + test('maps a recycled dirty session to a retryable public failure', () => { + const failure = publicExecutionFailure( + new Error('MICROVM_UNHEALTHY: Runtime session rt_private workspace was dirty and has been recycled'), + ); + expect(failure).toEqual({ + status: 503, + body: { + error: 'microvm_unhealthy', + message: 'Sandbox runtime is unavailable', + }, + }); + expect(JSON.stringify(failure)).not.toContain('rt_private'); + }); + + test('maps oversized input delivery to 413', () => { + const err = new Error('SESSION_INPUT_TOO_LARGE: Session inputs exceed the 536870912-byte budget'); + expect(publicExecutionFailure(err)).toEqual({ + status: 413, + body: { + error: 'session_input_too_large', + message: 'Input files exceed the delivery limit', + }, + }); + }); + + test('maps unavailable input objects to 422 without exposing object details', () => { + const failure = publicExecutionFailure( + new Error('SESSION_INPUT_UNAVAILABLE: Failed to fetch private/customer-list.csv from file-server.internal'), + ); + expect(failure).toEqual({ + status: 422, + body: { + error: 'session_input_unavailable', + message: 'One or more input files are unavailable', + }, + }); + expect(JSON.stringify(failure)).not.toContain('customer-list.csv'); + expect(JSON.stringify(failure)).not.toContain('file-server.internal'); + }); + + test('maps input source outages to 502 rather than MicroVM unavailability', () => { + const err = new Error('SESSION_INPUT_SOURCE_FAILED: upstream request failed with status 503'); + expect(publicExecutionFailure(err)).toEqual({ + status: 502, + body: { + error: 'session_input_source_failed', + message: 'Input file service is unavailable', + }, + }); + }); + + test('does not expose AWS identifiers embedded in backend failures', () => { + const arn = 'arn:aws:iam::123456789012:role/private-microvm-exec'; + const failure = publicExecutionFailure( + new Error(`MICROVM_LAUNCH_FAILED: Access denied while passing ${arn} to mvm-secret-123`), + ); + expect(failure).toEqual({ + status: 503, + body: { error: 'microvm_launch_failed', message: 'Sandbox launch failed' }, + }); + expect(JSON.stringify(failure)).not.toContain('123456789012'); + expect(JSON.stringify(failure)).not.toContain('mvm-secret-123'); + }); + test('maps sandbox request guard failures to public bad requests', () => { const axiosErr = { message: 'Request failed with status code 400', diff --git a/service/src/utils.ts b/service/src/utils.ts index 9c380cc..c117f2a 100644 --- a/service/src/utils.ts +++ b/service/src/utils.ts @@ -108,6 +108,51 @@ export function sandboxErrorMessageFromAxios(error: AxiosError): string { export function publicExecutionFailure(error: unknown): { status: number; body: { error: string; message: string } } | null { const message = error instanceof Error ? error.message : ''; + + /* Typed worker failures cross BullMQ as `: `. Runtime-session + * and MicroVM codes describe sandbox availability; SESSION_INPUT_* codes + * describe the caller's declared input set or its upstream object source. */ + const backendMatch = message.match( + /^(RUNTIME_SESSION_BUSY|MICROVM_[A-Z_]+|SESSION_INPUT_[A-Z_]+):\s*(.+)$/, + ); + if (backendMatch) { + const code = backendMatch[1]; + const statuses: Record = { + RUNTIME_SESSION_BUSY: 409, + SESSION_INPUT_TOO_LARGE: 413, + SESSION_INPUT_UNAVAILABLE: 422, + SESSION_INPUT_SOURCE_FAILED: 502, + SESSION_INPUT_PREPARATION_FAILED: 500, + SESSION_INPUT_ABORTED: 504, + }; + const sessionInputFailure = code.startsWith('SESSION_INPUT_'); + const status = statuses[code] ?? (sessionInputFailure ? 500 : 503); + const publicMessages: Record = { + RUNTIME_SESSION_BUSY: 'Runtime session is busy', + MICROVM_LAUNCH_FAILED: 'Sandbox launch failed', + MICROVM_LAUNCH_THROTTLED: 'Sandbox capacity is temporarily unavailable', + MICROVM_UNHEALTHY: 'Sandbox runtime is unavailable', + MICROVM_FENCED: 'Runtime session changed during execution', + MICROVM_DEADLINE_EXCEEDED: 'Sandbox execution deadline exceeded', + SESSION_INPUT_TOO_LARGE: 'Input files exceed the delivery limit', + SESSION_INPUT_UNAVAILABLE: 'One or more input files are unavailable', + SESSION_INPUT_SOURCE_FAILED: 'Input file service is unavailable', + SESSION_INPUT_PREPARATION_FAILED: 'Input files could not be prepared', + SESSION_INPUT_ABORTED: 'Input delivery timed out', + }; + /* The backend message is retained in worker/router logs, but it can contain + * AWS identifiers, internal endpoints, object ids, or file names. Only the + * stable code and a fixed public message cross the API boundary. */ + return { + status, + body: { + error: code.toLowerCase(), + message: publicMessages[code] + ?? (sessionInputFailure ? 'Input delivery failed' : 'Sandbox runtime is unavailable'), + }, + }; + } + const match = message.match(/^Error from sandbox(?:\s+\[([a-z_]+)\])?:\s*(?:\[([a-z_]+)\]\s*)?(.+)$/); if (!match) return null; diff --git a/service/src/workers.ts b/service/src/workers.ts index 177e8fd..769827d 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -3,26 +3,24 @@ import { Worker } from 'bullmq'; import type * as t from './types'; import { filterSystemLogs, applySystemReplacements, getAxiosErrorDetails, sandboxErrorMessageFromAxios } from './utils'; import { jobProcessingDuration, jobsCompleted, jobsFailed, activeJobs, workerRunning } from './metrics'; -import { Jobs, Queues } from './enum'; +import { Queues } from './enum'; import { connection } from './queue'; -import { env } from './config'; +import { env, jobDeadlineAtMs } from './config'; import { summarizeSandboxResponse, summarizeText } from './execution-log'; import { createGatewayEgressGrant, restoreGatewaySandboxResult, revokeGatewayEgressGrant } from './egress-gateway-client'; import { refreshEgressGrantClaims } from './sandbox-egress'; import { buildSandboxExecuteRequest } from './sandbox-dispatch'; +import { prepareInputDelivery } from './runtime-session/input-delivery'; +import { SessionFilesError } from './runtime-session/files'; +import { resolveRuntimeSessionIdForJob } from './runtime-session/job-policy'; +import { getSandboxBackend, SandboxBackendError, type SandboxRawResponse } from './sandbox-backend'; import { isSyntheticPrincipalSource } from './auth/synthetic'; -import { injectTraceHeaders, withSpan, withTraceContext } from './telemetry'; +import { withSpan, withTraceContext } from './telemetry'; import logger from './logger'; const { INSTANCE_ID } = env; const WORKER_ID = `${INSTANCE_ID}-${process.pid}`; -type SandboxLogResponse = t.ExecuteResponse & { - session_id: string; - files?: t.FileRefs; - run?: t.ExecuteResponse['run']; -}; - function isAbortError(error: unknown): boolean { return axios.isAxiosError(error) && (error.name === 'AbortError' || error.code === 'ERR_CANCELED'); } @@ -44,12 +42,20 @@ async function processJobInner(job: t.ExecuteJob): Promise { activeJobs.inc({ language }); const controller = new AbortController(); - const timer = setTimeout(() => controller.abort(), env.JOB_TIMEOUT); + const deadlineAtMs = jobDeadlineAtMs(job.timestamp, env.JOB_TIMEOUT); + const remainingBudgetMs = Math.max(0, deadlineAtMs - Date.now()); + const timer = remainingBudgetMs > 0 + ? setTimeout(() => controller.abort(), remainingBudgetMs) + : undefined; + if (remainingBudgetMs === 0) controller.abort(); let egressGrantId: string | undefined; let egressGrantTokenForRestore: string | undefined; let revokeReason = 'completed'; try { + if (controller.signal.aborted) { + throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); + } let sandboxPayload = payload; let executionManifestClaims = job.data.executionManifestClaims; let egressGrantToken = job.data.egressGrantToken; @@ -71,8 +77,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { : undefined; } + const delivery = prepareInputDelivery(payload, sandboxPayload); const sandboxRequest = buildSandboxExecuteRequest({ - payload: sandboxPayload, + payload: delivery.payload, egressGrantToken, executionManifestClaims, executionManifestPrivateKey: env.EXECUTION_MANIFEST_PRIVATE_KEY, @@ -81,32 +88,63 @@ async function processJobInner(job: t.ExecuteJob): Promise { }); egressGrantTokenForRestore = egressGrantToken; - const response = await withSpan('codeapi.sandbox.execute', { - 'http.request.method': 'POST', - 'url.path': `/${Jobs.execute}`, - 'codeapi.language': language, - }, () => axios.post( - `${env.SANDBOX_ENDPOINT}/${Jobs.execute}`, - sandboxRequest.body, - { - headers: injectTraceHeaders(sandboxRequest.headers), - signal: controller.signal, - } - ), 'CLIENT'); - - if (response.status !== 200) { - throw new Error('Error from sandbox'); - } + const runtimeSessionId = resolveRuntimeSessionIdForJob({ + mode: env.RUNTIME_SESSION_MODE, + runtimeSessionId: job.data.runtimeSessionId, + runtimeSessionExemption: job.data.runtimeSessionExemption, + isSynthetic: isSyntheticJob, + }); - const responseData = egressGrantTokenForRestore - ? await restoreGatewaySandboxResult({ + /* Stateful Lambda runs this inside its session commit barrier. The worker + * still invokes it unconditionally as the HTTP/stateless fallback; marking + * the transformed object makes that second call an idempotent no-op. */ + const resultRestoreToken = egressGrantTokenForRestore; + const finalizedSandboxResults = new WeakSet(); + const finalizeSandboxResult = async (result: SandboxRawResponse): Promise => { + if ( + resultRestoreToken === undefined || + resultRestoreToken.length === 0 || + finalizedSandboxResults.has(result) + ) { + return result; + } + const restored = await restoreGatewaySandboxResult({ grantId: egressGrantId, - egressGrantToken: egressGrantTokenForRestore, - result: response.data, + egressGrantToken: resultRestoreToken, + result, + isSynthetic: isSyntheticJob, + signal: controller.signal, + }); + finalizedSandboxResults.add(restored); + return restored; + }; + + const responseRaw = await getSandboxBackend().execute( + { + body: sandboxRequest.body, + headers: sandboxRequest.headers, + inputDelivery: delivery.refs, + }, + { + executionId: job.data.executionId ?? '', + language, isSynthetic: isSyntheticJob, signal: controller.signal, - }) - : response.data; + deadlineAtMs, + tenantId: job.data.tenantId, + canonicalUserId: job.data.canonicalUserId, + runtimeSessionId, + runtimeSessionMode: env.RUNTIME_SESSION_MODE, + /* Stateful backends run this as a commit barrier after user code but + * before checkpointing/reusing the mutated workspace. Stateless/HTTP + * paths retain the worker-owned fallback immediately below. */ + sessionResultFinalizer: resultRestoreToken !== undefined && resultRestoreToken.length > 0 + ? finalizeSandboxResult + : undefined, + }, + ); + + const responseData = await finalizeSandboxResult(responseRaw); if (!isSyntheticJob) { logger.info('Sandbox response', summarizeSandboxResponse(responseData)); @@ -150,11 +188,20 @@ async function processJobInner(job: t.ExecuteJob): Promise { return result; } catch (error) { - revokeReason = isAbortError(error) ? 'timeout' : 'failed'; + revokeReason = controller.signal.aborted || isAbortError(error) ? 'timeout' : 'failed'; const errorDetails = getAxiosErrorDetails(error); logger.error('Error processing job', errorDetails); - if (isAbortError(error)) { + if (controller.signal.aborted) { + throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); + } else if (error instanceof SandboxBackendError) { + throw new Error(`${error.code}: ${error.message}`); + } else if (error instanceof SessionFilesError) { + /* BullMQ serializes Error rather than preserving custom prototypes. + * Carry the stable code in the message so the public router can map the + * input failure without confusing it with MicroVM health. */ + throw new Error(`${error.code}: ${error.message}`); + } else if (isAbortError(error)) { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); } else if (axios.isAxiosError(error)) { /** Preserve error message from sandbox */ @@ -174,7 +221,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { logger.error('Failed to revoke egress grant', { grantId: egressGrantId, error: getAxiosErrorDetails(error) }); }); } - clearTimeout(timer); + if (timer) clearTimeout(timer); endTimer(); activeJobs.dec({ language }); }