diff --git a/api/Dockerfile b/api/Dockerfile index ff4e848..b3ee5b0 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -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/src/api/lifecycle.test.ts b/api/src/api/lifecycle.test.ts new file mode 100644 index 0000000..269b472 --- /dev/null +++ b/api/src/api/lifecycle.test.ts @@ -0,0 +1,46 @@ +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('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..2267efb --- /dev/null +++ b/api/src/api/lifecycle.ts @@ -0,0 +1,94 @@ +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; + + 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/v2.ts b/api/src/api/v2.ts index 5046ca3..53b7883 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -12,6 +12,12 @@ import { activeSandboxExecutions, recordSandboxExecution } from '../metrics'; import { classifySandboxSafeError } from '../safe-error'; import { withSpan } from '../telemetry'; import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; +import { + RUNTIME_SESSION_ID_HEADER, + bindSessionWorkspace, + parseSessionBindingFromHeader, +} from '../session-workspace'; +import { streamSessionCheckpoint, restoreSessionCheckpoint } from '../session-checkpoint'; const router = express.Router(); const SYNTHETIC_PRINCIPAL_SOURCE = 'synthetic_test'; @@ -130,6 +136,7 @@ function getJob( egressGrantToken?: string, toolCallSocketEnabled = false, isSynthetic = false, + runtimeSessionHeader?: string | string[], ): Job { const { session_id, language, version, args, stdin, files, @@ -171,6 +178,14 @@ 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); + const session = binding ? bindSessionWorkspace(binding) ?? null : null; + return new Job({ session_id: session_id ?? null, runtime: rt, @@ -194,6 +209,7 @@ function getJob( egress_grant: egressGrantToken, tool_call_socket_enabled: toolCallSocketEnabled, is_synthetic: isSynthetic, + session, }); } @@ -238,6 +254,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 streams a tar.gz body, not JSON. */ + if (req.path === '/session/restore') return next(); if (!req.headers['content-type']?.startsWith('application/json')) { return res.status(415).json({ message: 'requests must be of type application/json' }); } @@ -322,6 +340,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(); @@ -441,4 +460,38 @@ 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; + bindSessionWorkspace(binding); + return true; +} + +/* 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); +}); + export default router; diff --git a/api/src/config.ts b/api/src/config.ts index f8e2aba..863bca7 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 regardless of any /run payload. A VM additionally opts in + * per-launch via the /run runHookPayload. */ + 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( diff --git a/api/src/entrypoint.sh b/api/src/entrypoint.sh index 69145b1..bfdd8c6 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}" diff --git a/api/src/index.ts b/api/src/index.ts index 92c590d..b946813 100644 --- a/api/src/index.ts +++ b/api/src/index.ts @@ -7,6 +7,7 @@ import { initializeSandboxWorkspaceIsolation, startWorkspaceReaper } from './wor import { httpMetricsMiddleware, metricsHandler } from './metrics'; import { positiveInt, shutdownTelemetry, traceHttpRequest } from './telemetry'; import v2Router from './api/v2'; +import lifecycleRouter, { LIFECYCLE_HOOK_BASE_PATH } from './api/lifecycle'; const app = express(); @@ -29,6 +30,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) => { diff --git a/api/src/job.ts b/api/src/job.ts index 291d76d..ad4c5b7 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -28,6 +28,7 @@ import { type SandboxJobIdentity, type SandboxWorkspaceLease, } from './workspace-isolation'; +import type { SessionWorkspace } from './session-workspace'; import { DIRKEEP, SANDBOX_DIR_MODE, @@ -675,11 +676,20 @@ 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 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 +708,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 }); @@ -775,6 +788,42 @@ export class Job { * (inclusive) so the per-job outside UID can create siblings/children while * escaped sibling UIDs cannot traverse the workspace tree. */ + /** + * 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 async ensureDirNoFollow(target: string): Promise { + const rel = path.relative(this.submissionDir, target); + if (!rel || rel === '.') return; + if (rel === '..' || rel.startsWith('..' + path.sep)) { + throw new Error(`Workspace path escapes the sandbox: ${target}`); + } + let cursor = this.submissionDir; + 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. */ + await fsp.mkdir(cursor).catch((err: NodeJS.ErrnoException) => { + if (err.code !== 'EEXIST') throw err; + }); + continue; + } + if (st.isSymbolicLink()) { + throw new Error(`Refusing to prime through symlinked workspace path: ${path.relative(this.submissionDir, cursor)}`); + } + if (!st.isDirectory()) { + throw new Error(`Workspace ancestor is not a directory: ${path.relative(this.submissionDir, cursor)}`); + } + } + } + private async secureAncestors(leaf: string): Promise { const rel = path.relative(this.submissionDir, leaf); if (!rel || rel === '..' || rel.startsWith('..' + path.sep)) return; @@ -786,14 +835,28 @@ export class Job { * concurrently. Skip paths we've already chmodded to avoid N*M redundant * syscalls (N files × M shared ancestors). */ if (this.chmoddedDirs.has(cursor)) continue; + /* 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(this.submissionDir, cursor)}`); + } await applySandboxPathPermissions(cursor, this.sandboxIdentity(), SANDBOX_DIR_MODE); this.chmoddedDirs.add(cursor); } } async prime(): Promise { - this.jobIdentity = await acquireJobIdentity(this.log); - this.workspaceLease = await createSandboxWorkspace(this.jobIdentity); + 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 +866,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', ); @@ -815,7 +879,7 @@ export class Job { const fileOps: Promise[] = []; for (const file of this.files) { if (file.id) { - fileOps.push(this.downloadAndWriteFile(file).then(() => {})); + fileOps.push(this.primeInputFile(file)); } else if (file.content !== undefined) { fileOps.push(this.writeFile(file)); } @@ -823,6 +887,50 @@ export class Job { await Promise.all(fileOps); } + /** + * 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): Promise { + if (this.session && file.id && (await this.reusePrimedInput(file))) return; + const name = await this.downloadAndWriteFile(file); + if (this.session && file.id && name) { + /* 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, file.id, primed?.readOnly === true, primed?.hash); + } + } + + private async reusePrimedInput(file: TFile): Promise { + const session = this.session; + if (!session || !file.id) return false; + if (session.primedInputId(file.name) !== file.id) return false; + const filePath = path.join(this.submissionDir, file.name); + try { + const st = await fsp.lstat(filePath); + 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, + }); + return true; + } catch { + return false; + } + } + private fileEgressBaseUrl(): string { return config.egress_gateway_url || config.file_server_url; } @@ -963,8 +1071,15 @@ export class Job { validateFilePath(originalName, this.submissionDir); const finalPath = path.join(this.submissionDir, originalName); const finalParent = path.dirname(finalPath); - await fsp.mkdir(finalParent, { recursive: true }); + /* 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); + else await fsp.mkdir(finalParent, { recursive: true }); await this.secureAncestors(finalParent); + /* Clear a symlink/dir a prior session turn may have squatted at the + * target so the rename lands a fresh regular file in the workspace + * rather than following a link or failing on a directory. */ + if (this.session) await fsp.rm(finalPath, { force: true, recursive: true }); const hash = await this.streamToDisk(response, tempPath, finalPath); const readOnly = response.headers.get('x-read-only')?.toLowerCase() === 'true'; @@ -1057,13 +1172,24 @@ export class Job { const content = Buffer.from(file.content ?? '', (file.encoding as BufferEncoding) ?? 'utf8'); const parentDir = path.dirname(filePath); - await fsp.mkdir(parentDir, { recursive: true }); + if (this.session) await this.ensureDirNoFollow(parentDir); + else await fsp.mkdir(parentDir, { recursive: true }); await this.secureAncestors(parentDir); + /* 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); 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( @@ -1482,7 +1608,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 +1618,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 +1686,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 +1858,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 }, @@ -1818,6 +1991,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; diff --git a/api/src/session-checkpoint.test.ts b/api/src/session-checkpoint.test.ts new file mode 100644 index 0000000..28cec9b --- /dev/null +++ b/api/src/session-checkpoint.test.ts @@ -0,0 +1,38 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { restoreSessionCheckpoint, streamSessionCheckpoint } from './session-checkpoint'; +import { resetSessionWorkspaceStateForTests } from './session-workspace'; + +afterEach(resetSessionWorkspaceStateForTests); + +/** 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; +} + +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); + }); +}); diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts new file mode 100644 index 0000000..f0be84d --- /dev/null +++ b/api/src/session-checkpoint.ts @@ -0,0 +1,163 @@ +import { spawn } from 'child_process'; +import * as fsp from 'fs/promises'; +import * as path from 'path'; +import type { Request, Response } from 'express'; +import { pipeline } from 'stream/promises'; +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'; + +export class SessionCheckpointError extends Error {} + +/** Streams `tar -czf -` of the session workspace to the response. */ +export async function streamSessionCheckpoint(res: Response): Promise { + const session = getBoundSessionWorkspace(); + if (!session) { + res.status(409).json({ message: 'No session workspace is bound' }); + return; + } + const { dir } = await session.ownership(); + + /* Carry the priming/output-diff state into the archive so a relaunched VM + * rebuilds it (see restoreSessionCheckpoint). Written under the held session + * lock, so no concurrent user code sees it, and removed once tar has read it. + * The path is a collision point: + * - a prior exec's sandbox code can squat it as a symlink OR a directory + * (attack — remove it; `recursive` clears a dir else the write fails + * EISDIR, `force` ignores absence, and neither follows a symlink), or + * - user code can legitimately create a regular file with this name (their + * data — do NOT delete it; skip metadata persistence this turn so their + * file tars as normal workspace content). + * We only ever remove/restore a sidecar we actually wrote. */ + const metaPath = path.join(dir, SESSION_META_FILE); + const squat = await fsp.lstat(metaPath).catch(() => null); + let wroteSidecar = false; + if (squat?.isFile()) { + logger.warn('Session meta sidecar path holds a user file; skipping metadata persistence this checkpoint'); + } else { + await fsp.rm(metaPath, { force: true, recursive: true }); + await fsp.writeFile(metaPath, JSON.stringify({ marker: SESSION_META_MARKER, ...session.snapshotMeta() }), { flag: 'wx' }); + wroteSidecar = true; + } + + res.status(200); + res.setHeader('Content-Type', CHECKPOINT_CONTENT_TYPE); + const tar = spawn('tar', ['-czf', '-', '-C', SANDBOX_WORKSPACE_ROOT, SESSION_WORKSPACE_ID], { + stdio: ['ignore', 'pipe', 'pipe'], + }); + tar.stderr.on('data', (chunk: Buffer) => logger.debug({ tar: chunk.toString() }, 'checkpoint tar')); + try { + await pipeline(tar.stdout, res); + const code: number = await new Promise((resolve) => tar.on('close', resolve)); + if (code !== 0) throw new SessionCheckpointError(`checkpoint tar exited ${code}`); + } 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 (wroteSidecar) await fsp.rm(metaPath, { force: true }).catch(() => {}); + } +} + +/** Extracts a `tar.gz` from the request body into the session workspace and + * re-owns it to the session's pinned UID. */ +export async function restoreSessionCheckpoint(req: Request, res: Response): Promise { + const session = getBoundSessionWorkspace(); + if (!session) { + res.status(409).json({ message: 'No session workspace is bound' }); + return; + } + const { dir, uid, gid } = await session.ownership(); + + /* Start from a clean workspace so a restore is a full replace, not a merge. */ + await fsp.rm(dir, { recursive: true, force: true }); + await fsp.mkdir(dir, { recursive: true }); + + /* Archives are created with relative `session/...` members (see the create + * side). Strip that leading component and extract straight into the session + * `dir`, so a poisoned archive member (`../x`, `ws_other/x`, `rootfile`) can + * never escape the workspace into shared runner space — every member lands + * under `dir`. Production hardening: verify/scan the archive before trusting a + * restore from shared storage. */ + const tar = spawn('tar', ['-xzf', '-', '--strip-components=1', '-C', dir], { + stdio: ['pipe', 'ignore', 'pipe'], + }); + tar.stderr.on('data', (chunk: Buffer) => logger.debug({ tar: chunk.toString() }, 'restore tar')); + try { + await pipeline(req, tar.stdin); + const code: number = await new Promise((resolve) => tar.on('close', resolve)); + if (code !== 0) throw new SessionCheckpointError(`restore tar exited ${code}`); + await chownRecursive(dir, uid, gid); + await applyRestoredMeta(session, dir); + res.status(200).json({ status: 'restored', dir: path.basename(dir) }); + } catch (error) { + logger.error({ err: error }, 'Failed to restore session checkpoint'); + /* A corrupt archive or cut-off upload can leave partially-extracted members + * behind. The control plane treats restore failure as non-fatal and runs the + * job anyway, so wipe the workspace to a clean slate — otherwise the job runs + * against a mix of stale checkpoint files instead of an empty workspace. */ + await fsp.rm(dir, { recursive: true, force: true }).catch(() => {}); + await fsp.mkdir(dir, { recursive: true }).catch(() => {}); + if (!res.headersSent) res.status(500).json({ message: 'restore failed' }); + } +} + +/** Applies the restored priming/output-diff sidecar to the bound session and + * removes it from disk so user code never sees it. Absent (older checkpoint) + * or malformed metadata is non-fatal — the session just re-primes/re-surfaces + * as if warmth were lost. */ +async function applyRestoredMeta(session: SessionWorkspace, dir: string): Promise { + const metaPath = path.join(dir, SESSION_META_FILE); + try { + /* Only trust a regular file: a restored archive is untrusted, so never + * follow a symlinked sidecar (it would read an arbitrary file into the + * loaded metadata). lstat does not follow the link. */ + const stat = await fsp.lstat(metaPath).catch(() => null); + if (!stat?.isFile()) return; + const parsed = JSON.parse(await fsp.readFile(metaPath, 'utf8')) as SessionMetaSnapshot; + /* Require our authenticity marker: a user file sharing the reserved name is + * never loaded as metadata nor deleted, even if it happens to contain + * primed/surfaced arrays. */ + if (parsed?.marker === SESSION_META_MARKER + && Array.isArray(parsed?.primed) && Array.isArray(parsed?.surfaced)) { + session.loadMeta(parsed); + await fsp.rm(metaPath, { force: true }).catch(() => {}); + } + } catch (error) { + logger.debug({ err: error }, 'No session meta sidecar to restore'); + } +} + +async function chownRecursive(dir: string, uid: number, gid: number): Promise { + await fsp.lchown(dir, uid, gid).catch(() => {}); + 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. `lchown` the link itself and never recurse + * through it (Dirent reports the link type, so `isDirectory()` is false). */ + if (entry.isSymbolicLink()) { + await fsp.lchown(full, uid, gid).catch(() => {}); + continue; + } + await fsp.chown(full, uid, gid).catch(() => {}); + if (entry.isDirectory()) await chownRecursive(full, uid, gid); + } +} diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts new file mode 100644 index 0000000..096c772 --- /dev/null +++ b/api/src/session-workspace.test.ts @@ -0,0 +1,156 @@ +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('')).toBeUndefined(); + expect(parseSessionBindingFromHeader(['rt_a', 'rt_b'])).toBeUndefined(); + expect(parseSessionBindingFromHeader('rt bad space')).toBeUndefined(); + expect(parseSessionBindingFromHeader('a'.repeat(129))).toBeUndefined(); + }); + + 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); + }); + + test('a different runtime session replaces the binding', () => { + const a = bindSessionWorkspace({ runtimeSessionId: 'rt_1' }); + const b = bindSessionWorkspace({ runtimeSessionId: 'rt_2' }); + expect(b).not.toBe(a); + expect(getBoundSessionWorkspace()?.runtimeSessionId).toBe('rt_2'); + }); + + 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(); + }); +}); diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts new file mode 100644 index 0000000..fa0dce6 --- /dev/null +++ b/api/src/session-workspace.ts @@ -0,0 +1,249 @@ +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'; + +/** Sidecar file the checkpoint tar carries so a relaunched VM rebuilds the + * in-memory priming/output-diff state a warm VM would have kept. Written into + * the workspace only while the session lock is held (no concurrent user code), + * and removed from disk again before any execute runs. */ +export const SESSION_META_FILE = '.codeapi-session-meta.json'; + +/** Authenticity marker embedded in the sidecar so restore can tell OUR metadata + * from a user file that merely shares the reserved name and happens to contain + * primed/surfaced arrays — without it, restore would loadMeta (poisoning the + * priming maps) and delete a legitimate user file. */ +export const SESSION_META_MARKER = 'codeapi.session-meta.v1'; + +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; +} + +/** 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 (typeof headerValue !== 'string') return undefined; + const runtimeSessionId = headerValue.trim(); + if (!RUNTIME_SESSION_ID_PATTERN.test(runtimeSessionId)) { + if (runtimeSessionId.length > 0) logger.warn('Ignoring malformed X-Runtime-Session-Id header'); + return undefined; + } + 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(); + + 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 }); + } + + /** 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 the checkpoint so a + * relaunched VM restores it (see {@link SESSION_META_FILE}). */ + snapshotMeta(): SessionMetaSnapshot { + return { + primed: [...this.primed.entries()], + surfaced: [...this.surfaced.entries()], + }; + } + + /** Rebuilds priming + output-diff state from a restored checkpoint sidecar. + * 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 { + 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. */ + async reset(): Promise { + await resetSessionWorkspace(); + this.surfaced.clear(); + this.primed.clear(); + this.lease = undefined; + if (this.identity) { + sandboxJobUidPool.release(this.identity); + this.identity = undefined; + } + } +} + +let boundSession: SessionWorkspace | undefined; + +/** Called by the `/run` lifecycle hook. Binding the same session twice is a + * no-op; a different runtime session id resets the prior one first. */ +export function bindSessionWorkspace(binding: SessionBinding | undefined): SessionWorkspace | undefined { + if (!binding) return boundSession; + if (boundSession && boundSession.runtimeSessionId === binding.runtimeSessionId) { + return boundSession; + } + if (boundSession) { + void boundSession.reset().catch((err) => logger.error({ err }, 'Failed to reset superseded session workspace')); + } + 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/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..a9d47f9 100644 --- a/api/src/workspace-isolation.ts +++ b/api/src/workspace-isolation.ts @@ -16,6 +16,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 +366,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); diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md new file mode 100644 index 0000000..c57ee15 --- /dev/null +++ b/docs/lambda-microvm/README.md @@ -0,0 +1,346 @@ +# 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). +- Design deep-dive → [`../aws-lambda-microvm-stateful-sessions-report.md`](../aws-lambda-microvm-stateful-sessions-report.md). + +> 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). The VM can die at any time; the session survives. + +--- + +## 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.5 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). + +--- + +## Quick start (from scratch) + +### 1. Provision AWS prerequisites (Terraform) + +```bash +cd docs/lambda-microvm/terraform +cp terraform.tfvars.example terraform.tfvars # edit region, name_prefix, image_name +terraform init && terraform apply +``` + +This creates the checkpoint bucket (encrypted, versioned, lifecycle-expired), an +artifact bucket, the **build role** (trust includes `sts:TagSession`; perms +include `logs:*` + `s3:GetObject` + optional ECR pull), a **logging-only +execution role**, and the build + runtime CloudWatch log groups. Capture the +outputs: + +```bash +terraform output # build_role_arn, execution_role_arn, checkpoint_bucket, artifact_bucket, ... +``` + +### 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=.dkr.ecr.us-east-1.amazonaws.com/codeapi-microvm-runner +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`). + +### 3. Create the MicroVM image (hookless) + +```bash +cd service +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 '{"FILE_SERVER_URL":"https://files.internal","EGRESS_GATEWAY_URL":"https://egress.internal"}' +# → prints LAMBDA_MICROVM_IMAGE_ARN when CREATED (~3-5 min) +``` + +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 +file-server / egress / manifest config via `--env-json` (or `MICROVM_IMAGE_ENV_JSON`) +— typically `FILE_SERVER_URL`, `EGRESS_GATEWAY_URL`, +`SANDBOX_ALLOWED_LOCAL_NETWORK_PORT`, `SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY`, +`SANDBOX_REQUIRE_EGRESS_MANIFEST`, `REQUIRE_EXECUTION_MANIFEST`. Without the +file/egress URLs the runner builds invalid `/sessions/...` URLs and can't fetch +inputs or upload outputs. To ship new runner code later, re-run +`build … push zip upload` and call the helper with `--update`. + +### 4. Configure the CodeAPI service + +```bash +CODEAPI_SANDBOX_BACKEND=lambda-microvm +CODEAPI_RUNTIME_SESSION_MODE=affinity # warm sessions + checkpoints +LAMBDA_MICROVM_IMAGE_ARN= +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=arn:aws:lambda:us-east-1:aws:network-connector:aws-network-connector:INTERNET_EGRESS + +# checkpoints (S3-compatible, same client as file-server) +CODEAPI_CHECKPOINT_BUCKET= +MINIO_ENDPOINT=s3.us-east-1.amazonaws.com +MINIO_PORT=443 # required: the client defaults to 9000, which fails against S3 +MINIO_USE_SSL=true +MINIO_REGION=us-east-1 +MINIO_ACCESS_KEY=... # from your CodeAPI task role, or the optional TF IAM user +MINIO_SECRET_KEY=... +``` + +`PTC_MODE` must be `replay` (the default) or unset — see +[Programmatic Tool Calling](#programmatic-tool-calling-ptc). + +### 5. 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 + +All names as they appear in `service/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 an execution waits for the session lock before falling back (affinity) or erroring (strict). | + +### MicroVM launch + +| Env | Default | Meaning | +|---|---|---| +| `LAMBDA_MICROVM_IMAGE_ARN` | — (required) | The image created in step 3. | +| `LAMBDA_MICROVM_IMAGE_VERSION` | latest | Pin a specific image version. | +| `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` | `300` | idlePolicy: auto-suspend after this idle. | +| `LAMBDA_MICROVM_SUSPEND_SECONDS` | `1800` | idlePolicy: auto-terminate after this suspended. | +| `LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS` | `300` | Proxy auth token TTL (cached to 80%). | +| `LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS` | `60000` | Budget for RunMicrovm → RUNNING. | +| `LAMBDA_MICROVM_HEALTH_TIMEOUT_MS` | `5000` | Health check budget. | +| `LAMBDA_MICROVM_LAUNCH_TPS` / `_RESUME_TPS` / `_SUSPEND_TPS` | `4` / `4` / `1` | Client-side throttle (headroom under AWS's 5/5/2 caps). | +| `LAMBDA_MICROVM_ALLOW_SHELL` | `false` | Must stay false in prod (shell auth token → IAM-deny). | + +### Checkpoints (affinity/strict only) + +| Env | Default | Meaning | +|---|---|---| +| `CODEAPI_SESSION_CHECKPOINTS` | `true` | `false` disables checkpoint/restore (sessions still reuse a warm workspace, but expiry recovery falls back to file-refs). | +| `CODEAPI_CHECKPOINT_BUCKET` | `MINIO_BUCKET` | Checkpoint bucket. | +| `CODEAPI_CHECKPOINT_PREFIX` | `rtsx-checkpoints/` | Key prefix. Objects are `.tar.gz`. | +| `CODEAPI_CHECKPOINT_MAX_BYTES` | `536870912` | Max checkpoint size (512 MiB). | +| `CODEAPI_CHECKPOINT_TIMEOUT_MS` | `60000` | Checkpoint transfer budget. | +| `MINIO_ENDPOINT` / `_PORT` / `_USE_SSL` / `_ACCESS_KEY` / `_SECRET_KEY` / `_REGION` / `_SESSION_TOKEN` | — | S3-compatible client (shared with file-server). Point at real S3 in prod. | + +--- + +## 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`. If the + session lock is contended past `LOCK_WAIT_MS`, fall back to a correct + stateless one-shot (warmth is only an optimization; the payload still carries + file refs). This is the recommended default for stateful sessions. +- **`strict`** — same, but lock contention returns HTTP 409 instead of falling + back. Use when you require a single serialized session and would rather fail + than run cold. + +--- + +## 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.** 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 and image builds are reliable. If you +later need build/runtime lifecycle hooks (e.g. an exact suspend-time checkpoint +flush), rebase the `lambda-microvm-runner` target on the snapshot-compatible +Lambda base container image (`public.ecr.aws/lambda/microvms:al2023-minimal`), +which bakes the hook-routing service components and a snapshot-safe OpenSSL. Only +then do hooks route. See [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 +`s3..amazonaws.com`, `MINIO_PORT=443`, `MINIO_USE_SSL=true`). The client +authenticates with static `MINIO_ACCESS_KEY`/`MINIO_SECRET_KEY` only — it does +**not** yet load task-role/IRSA credentials — so set +`create_checkpoint_access_user = true` and use its keys. Attaching +`checkpoint_access_policy_arn` to a task role alone will not work until the +checkpoint client grows IRSA-aware credential loading (a tracked follow-up). + +**Egress posture.** For dev, the Lambda-managed `INTERNET_EGRESS` connector gives +default public egress. For hardened prod, set +`CODEAPI_HARDENED_SANDBOX_MODE=true` and provide a VPC egress connector + SG +locked to your egress-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. +- **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 perms must include `logs:*` + `s3:GetObject` (+ ECR 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. + +--- + +## Teardown + +```bash +# Terminate only VMs launched from THIS image, then delete the image. +# IMAGE_ARN scopes the sweep so it never touches unrelated MicroVMs in a shared +# account (ListMicrovms returns every VM in the region). +cd service +export IMAGE_ARN="arn:aws:lambda:us-east-1::microvm-image:codeapi-session" +AWS_PROFILE=... bun -e 'import { LambdaMicrovmsClient, ListMicrovmsCommand, TerminateMicrovmCommand, DeleteMicrovmImageCommand } from "@aws-sdk/client-lambda-microvms"; const arn=process.env.IMAGE_ARN; const c=new LambdaMicrovmsClient({region:"us-east-1"}); const v=await c.send(new ListMicrovmsCommand({})) as any; for (const m of (v.microvms??[]).filter((x:any)=>!/TERMINAT/.test(x.state) && x.imageArn===arn)) await c.send(new TerminateMicrovmCommand({microvmIdentifier:m.microvmId})); await c.send(new DeleteMicrovmImageCommand({imageIdentifier:"codeapi-session"})).catch(()=>{});' + +# then the prerequisites +cd ../docs/lambda-microvm/terraform && 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..1b05614 --- /dev/null +++ b/docs/lambda-microvm/terraform/.gitignore @@ -0,0 +1,6 @@ +.terraform/ +.terraform.lock.hcl +*.tfstate +*.tfstate.* +terraform.tfvars +crash.log diff --git a/docs/lambda-microvm/terraform/README.md b/docs/lambda-microvm/terraform/README.md new file mode 100644 index 0000000..abde0b6 --- /dev/null +++ b/docs/lambda-microvm/terraform/README.md @@ -0,0 +1,44 @@ +# 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-KMS), versioned, public-access + blocked, lifecycle-expired (`checkpoint_retention_days`). +- **Artifact S3 bucket** — for the code-artifact zip (optional; reuse an existing + one with `create_artifact_bucket = false` + `artifact_bucket_name`). +- **Build role** — assumed by Lambda during `create/update-microvm-image`. Trust + includes `sts:TagSession`; perms include `logs:*`, `s3:GetObject` on the + artifact bucket, and (optional) private-ECR pull. Getting this wrong yields a + build failure with an empty `stateReason`. +- **Execution role** — logging-only least-privilege, for `RunMicrovm`. +- **CloudWatch log groups** — build (`/aws/lambda-microvms/`) and + runtime. +- **Checkpoint access** — an IAM policy to attach to your CodeAPI task role + (preferred), or an optional IAM user + access key + (`create_checkpoint_access_user = true`) for non-role deployments. + +## Usage + +```bash +cp terraform.tfvars.example terraform.tfvars # edit +terraform init +terraform apply +terraform output +``` + +## Notes + +- 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. +- `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`. Prefer the task-role policy when you can. +- Buckets use `force_destroy = true` so `terraform destroy` is clean in dev. + Reconsider for prod. diff --git a/docs/lambda-microvm/terraform/main.tf b/docs/lambda-microvm/terraform/main.tf new file mode 100644 index 0000000..34abdf6 --- /dev/null +++ b/docs/lambda-microvm/terraform/main.tf @@ -0,0 +1,276 @@ +data "aws_caller_identity" "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 + + base_tags = merge(var.tags, { + "app" = "codeapi" + "component" = "lambda-microvm" + }) +} + +# -------------------------------------------------------------------------- +# 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 = true + 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" + } + } +} + +# -------------------------------------------------------------------------- +# 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 = true + 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 + rule { + id = "expire-checkpoints" + status = "Enabled" + filter {} + expiration { + days = var.checkpoint_retention_days + } + noncurrent_version_expiration { + noncurrent_days = var.checkpoint_retention_days + } + } +} + +# -------------------------------------------------------------------------- +# 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 logs:* + 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:aws:s3:::${local.artifact_bucket}/*"] + } + + statement { + sid = "BuildLogs" + effect = "Allow" + actions = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"] + # Two ARN forms: the log-group itself (CreateLogGroup) and its streams + # (`:*` suffix — CreateLogStream/PutLogEvents act at stream level, which the + # `/*` form does NOT match). Missing `:*` = build fails with empty stateReason. + resources = [ + "arn:aws:logs:${var.region}:${local.account_id}:log-group:/aws/lambda-microvms/*", + "arn:aws:logs:${var.region}:${local.account_id}:log-group:/aws/lambda-microvms/*:*", + ] + } + + dynamic "statement" { + for_each = var.private_ecr ? [1] : [] + content { + sid = "PrivateEcrPull" + effect = "Allow" + actions = [ + "ecr:GetAuthorizationToken", + "ecr:BatchCheckLayerAvailability", + "ecr:GetDownloadUrlForLayer", + "ecr:BatchGetImage", + ] + resources = ["*"] + } + } +} + +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:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"] + # Group ARN for CreateLogGroup + the `:*` stream form for the stream-level + # actions (the group is pre-created here, but grant both so the runtime + # logging agent works regardless). + resources = [ + aws_cloudwatch_log_group.runtime.arn, + "${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 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_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..3f6dfee --- /dev/null +++ b/docs/lambda-microvm/terraform/outputs.tf @@ -0,0 +1,53 @@ +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 "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 as --build-role-arn." + 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 = <<-EOT + IAM policy for checkpoint S3 access. NOTE: the checkpoint client reads static + MINIO_ACCESS_KEY/SECRET and does not yet load task-role/IRSA credentials, so + attaching this to a task role alone does not grant access today — use + create_checkpoint_access_user (static keys). This output is for when + IRSA-capable credential loading lands in the checkpoint client. + EOT + value = aws_iam_policy.checkpoint_access.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..a2e663c --- /dev/null +++ b/docs/lambda-microvm/terraform/terraform.tfvars.example @@ -0,0 +1,25 @@ +# 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 + +# Checkpoints are a resumable cache; keep the window short. +checkpoint_retention_days = 14 +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..4258e96 --- /dev/null +++ b/docs/lambda-microvm/terraform/variables.tf @@ -0,0 +1,90 @@ +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 "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 +} + +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_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). This is currently the ONLY working path: the checkpoint + client reads static keys and does not load task-role/IRSA credentials, so + attaching `checkpoint_access_policy_arn` to a task role alone does not work + yet. Set true unless you supply MINIO_ACCESS_KEY/SECRET some other way. + 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..26471d8 --- /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" + } + } +} + +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..b6d94c3 --- /dev/null +++ b/scripts/build-lambda-microvm-artifact.sh @@ -0,0 +1,108 @@ +#!/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 +# 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")/.." + +IMAGE_TAG="${IMAGE_TAG:-$(git rev-parse --short HEAD)}" +ECR_URI="${ECR_URI:-}" +S3_URI="${S3_URI:-}" +OUT_DIR="${OUT_DIR:-.build-lambda-microvm}" +LOCAL_TAG="codeapi-lambda-microvm-runner:${IMAGE_TAG}" + +require_ecr() { + [ -n "$ECR_URI" ] || { echo "ECR_URI is required for this stage" >&2; exit 1; } +} + +do_build() { + echo ">> buildx arm64 lambda-microvm-runner (${LOCAL_TAG})" + docker buildx build \ + --platform linux/arm64 \ + --target lambda-microvm-runner \ + -f api/Dockerfile \ + -t "$LOCAL_TAG" \ + ${ECR_URI:+-t "$ECR_URI:$IMAGE_TAG"} \ + --load \ + . +} + +do_push() { + require_ecr + 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%%/*}" + docker push "$ECR_URI:$IMAGE_TAG" +} + +do_zip() { + require_ecr + mkdir -p "$OUT_DIR" + cat > "$OUT_DIR/Dockerfile" <> wrote $OUT_DIR/artifact.zip (FROM ${ECR_URI}:${IMAGE_TAG})" +} + +do_upload() { + [ -n "$S3_URI" ] || { echo "S3_URI is required for upload" >&2; exit 1; } + 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 < \\ + --additional-os-capabilities '["ALL"]' \\ + --region \${AWS_REGION} + +Notes: +- env vars (SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY, SANDBOX_REQUIRE_EGRESS_MANIFEST, + EGRESS_GATEWAY_URL, SANDBOX_ALLOWED_LOCAL_NETWORK_PORT) are image-build-time + config: pass them via --environment-variables on create/update-microvm-image. +- 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..d8beae8 100644 --- a/service/bun.lock +++ b/service/bun.lock @@ -5,6 +5,7 @@ "": { "name": "code-execution-lambda", "dependencies": { + "@aws-sdk/client-lambda-microvms": "3.1079.0", "@opentelemetry/api": "1.9.1", "@opentelemetry/core": "2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "0.218.0", @@ -50,6 +51,38 @@ }, }, "packages": { + "@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/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/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 +177,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 +361,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=="], diff --git a/service/package.json b/service/package.json index 89ca1ee..d02b0f2 100644 --- a/service/package.json +++ b/service/package.json @@ -26,6 +26,7 @@ "license": "Apache-2.0", "description": "", "dependencies": { + "@aws-sdk/client-lambda-microvms": "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.ts b/service/scripts/create-microvm-image.ts new file mode 100644 index 0000000..bef3bbb --- /dev/null +++ b/service/scripts/create-microvm-image.ts @@ -0,0 +1,136 @@ +/** + * 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 + * --region MICROVM_REGION default us-east-1 + * --memory MICROVM_MEMORY_MIB baseline memory (default 2048) + * --update MICROVM_UPDATE=true update an existing image (new version) instead of create + */ +import { + LambdaMicrovmsClient, + CreateMicrovmImageCommand, + UpdateMicrovmImageCommand, + GetMicrovmImageCommand, +} from '@aws-sdk/client-lambda-microvms'; + +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 memory = Number(arg('--memory', 'MICROVM_MEMORY_MIB', '2048')); +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 file-server / egress-gateway / manifest config HERE or + * it builds invalid `/sessions/...` URLs and can't fetch inputs or upload + * outputs. The helper can't know your deployment's URLs, so pass them via + * --env-json '{"FILE_SERVER_URL":"...","EGRESS_GATEWAY_URL":"...", ...}' (or the + * MICROVM_IMAGE_ENV_JSON env). Typical keys: FILE_SERVER_URL, EGRESS_GATEWAY_URL, + * SANDBOX_ALLOWED_LOCAL_NETWORK_PORT, SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY, + * SANDBOX_REQUIRE_EGRESS_MANIFEST, REQUIRE_EXECUTION_MANIFEST. */ +function parseEnvJson(): Record { + const raw = arg('--env-json', 'MICROVM_IMAGE_ENV_JSON'); + if (!raw) return {}; + const parsed = JSON.parse(raw) as Record; + return Object.fromEntries(Object.entries(parsed).map(([k, v]) => [k, String(v)])); +} + +/* 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, + 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}...`); + if (isUpdate) { + await client.send(new UpdateMicrovmImageCommand({ imageIdentifier: name, ...shared } as never)); + } else { + await client.send(new CreateMicrovmImageCommand({ name, description: 'CodeAPI hookless session runner', ...shared } as never)); + } + + const started = Date.now(); + /* Cap the wait so a build wedged in CREATING/UPDATING can't hang a + * provisioning job forever (observed during the spike). Override with + * MICROVM_BUILD_DEADLINE_MINUTES. */ + const deadlineMs = started + Number(process.env.MICROVM_BUILD_DEADLINE_MINUTES ?? '30') * 60_000; + for (;;) { + /* GetMicrovmImage accepts the image name as the identifier. */ + const img = (await client.send( + new GetMicrovmImageCommand({ imageIdentifier: name }), + )) as { state?: string; imageArn?: string; imageVersion?: string; stateReason?: string }; + const elapsed = Math.round((Date.now() - started) / 1000); + if (Date.now() > deadlineMs) { + console.error(`\nTimed out after ${elapsed}s still in state ${img.state ?? 'UNKNOWN'}. Check the build log group.`); + process.exit(1); + } + const state = img.state ?? 'UNKNOWN'; + if (state === 'CREATED' || state === 'UPDATED') { + console.log(`\n${state} in ${elapsed}s`); + console.log(` imageArn: ${img.imageArn}`); + console.log(` version: ${img.imageVersion ?? '(latest)'}`); + console.log('\nSet on the CodeAPI service:'); + console.log(` LAMBDA_MICROVM_IMAGE_ARN=${img.imageArn}`); + return; + } + if (state.includes('FAILED')) { + console.error(`\n${state} after ${elapsed}s. reason: ${img.stateReason || '(empty — check the build log group)'}`); + process.exit(1); + } + if (elapsed % 60 < 20) console.log(` [${elapsed}s] ${state}`); + await new Promise((r) => setTimeout(r, 20_000)); + } +} + +main().catch((error) => { + console.error('create-microvm-image failed:', error instanceof Error ? error.message : error); + process.exit(1); +}); diff --git a/service/src/config.ts b/service/src/config.ts index 5abc57c..5f0c30e 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -49,6 +49,12 @@ 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; +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 function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, jobTimeoutMs: number): number { const defaultTtlSeconds = Math.max(1, Math.ceil((jobTimeoutMs + EGRESS_GRANT_GRACE_MS) / 1000)); if (rawTtlSeconds == null || rawTtlSeconds.trim() === '') { @@ -142,6 +148,64 @@ 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. Startup policy rejects + * this value until the backend implementation lands. + */ + 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`: best-effort session reuse; contention falls back to a + * stateless one-shot execution (correct because payloads always carry + * file refs — warmth is only an optimization). + * - `strict`: sessions required; contention surfaces as HTTP 409. + */ + 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: Number(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), + LAMBDA_MICROVM_PORT: Number(process.env.LAMBDA_MICROVM_PORT) || 8080, + LAMBDA_MICROVM_MAX_DURATION_SECONDS: Math.min( + Number(process.env.LAMBDA_MICROVM_MAX_DURATION_SECONDS) || 28_800, + 28_800, + ), + LAMBDA_MICROVM_IDLE_SECONDS: Number(process.env.LAMBDA_MICROVM_IDLE_SECONDS) || 300, + LAMBDA_MICROVM_SUSPEND_SECONDS: Number(process.env.LAMBDA_MICROVM_SUSPEND_SECONDS) || 1_800, + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: Number(process.env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS) || 300, + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: Number(process.env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS) || 60_000, + LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: Number(process.env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS) || 5_000, + LAMBDA_MICROVM_LAUNCH_TPS: Number(process.env.LAMBDA_MICROVM_LAUNCH_TPS) || 4, + LAMBDA_MICROVM_RESUME_TPS: Number(process.env.LAMBDA_MICROVM_RESUME_TPS) || 4, + LAMBDA_MICROVM_SUSPEND_TPS: Number(process.env.LAMBDA_MICROVM_SUSPEND_TPS) || 1, + /* 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_TOKEN_TPS: Number(process.env.LAMBDA_MICROVM_TOKEN_TPS) || 8, + 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: Number(process.env.CODEAPI_CHECKPOINT_MAX_BYTES) || 512 * 1024 * 1024, + CHECKPOINT_TIMEOUT_MS: Number(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/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..0175cca 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -194,6 +194,67 @@ 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 runtimeSessionFallback = new Counter({ + name: 'codeapi_runtime_session_fallback_total', + help: 'Affinity-mode executions that fell back to a stateless one-shot VM', +}); + +export const microvmActiveSessions = new Gauge({ + name: 'codeapi_microvm_active_sessions', + help: 'Live runtime sessions tracked in the registry active set', +}); + +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..3554419 --- /dev/null +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, test } from 'bun:test'; +import { + MemoryCheckpointStore, + CheckpointTooLargeError, + checkpointObjectKey, + checkpointPrefixFor, +} from './checkpoint-store'; + +const BIG = 1_000_000; + +describe('checkpoint store', () => { + 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('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); + + const fetched = await store.get('rt_1', BIG); + expect(fetched?.toString()).toBe('workspace-bytes'); + /* stored copy is independent of the caller's buffer */ + original.fill(0); + expect((await store.get('rt_1', BIG))?.toString()).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')); + /* 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 store.get('rt_1', BIG))?.toString()).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('put prunes strictly-older sequences for the session', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_1', 1, Buffer.from('v1')); + await store.put('rt_1', 2, Buffer.from('v2')); + await store.put('rt_other', 1, Buffer.from('other')); + 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_other', 1)); + }); + + test('rejects a checkpoint larger than maxBytes', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_big', 1, Buffer.alloc(2048)); + 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..957c7ac --- /dev/null +++ b/service/src/runtime-session/checkpoint-store.ts @@ -0,0 +1,175 @@ +import { Client as MinioClient } from 'minio'; +import { env } from '../config'; + +/** + * Durable storage for session workspace checkpoints. Each checkpoint writes a + * distinct object under a per-session prefix keyed by a monotonic sequence + * (`/.tar.gz`); restore reads the + * lexicographically-greatest (highest-sequence) object. The sequence comes from + * a Redis INCR (see `allocateCheckpointSequence`) that the caller seeds above + * `latestSequence()` if the counter was lost to a TTL — so ordering is a pure + * server-side monotonic counter with NO dependence on per-worker wall clocks + * (which can skew or step backwards across pods) and no reset below retained + * objects. A relaunch finds the latest by listing the prefix even with no + * registry record. Older objects are best-effort pruned after each put. + */ +export interface CheckpointStore { + put(runtimeSessionId: string, sequence: number, data: Buffer): Promise; + /** Reads the highest-sequence checkpoint for the session. `maxBytes` is + * enforced BEFORE the object is buffered into memory (stat the object first + * for S3-backed stores) so a stray oversized checkpoint can't OOM the worker. + * Throws {@link CheckpointTooLargeError} when exceeded. */ + get(runtimeSessionId: string, maxBytes: number): 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 {} + +/** 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`; +} + +/** 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: MinioClient; + private readonly bucket: string; + + constructor() { + this.client = new MinioClient({ + endPoint: process.env.MINIO_ENDPOINT ?? 'localhost', + port: Number(process.env.MINIO_PORT) || 9000, + useSSL: process.env.MINIO_USE_SSL === 'true', + accessKey: process.env.MINIO_ACCESS_KEY ?? '', + secretKey: process.env.MINIO_SECRET_KEY ?? '', + ...(process.env.MINIO_SESSION_TOKEN ? { sessionToken: process.env.MINIO_SESSION_TOKEN } : {}), + ...(process.env.MINIO_REGION ? { region: process.env.MINIO_REGION } : {}), + }); + /* `||` 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 = process.env.CODEAPI_CHECKPOINT_BUCKET || process.env.MINIO_BUCKET || 'test-bucket'; + } + + async put(runtimeSessionId: string, sequence: number, data: Buffer): Promise { + const key = checkpointObjectKey(runtimeSessionId, sequence); + await this.client.putObject(this.bucket, key, data, data.length, { + 'Content-Type': 'application/x-gtar', + }); + /* Best-effort: drop only STRICTLY-OLDER sequences. Never touch a newer key + * (a crash-orphaned late put must not delete the checkpoint that superseded + * it — restore always reads the max sequence). */ + await this.pruneOlderThan(runtimeSessionId, key).catch(() => {}); + } + + async latestSequence(runtimeSessionId: string): Promise { + const key = await this.latestKey(runtimeSessionId); + return key ? sequenceFromKey(key) : 0; + } + + async get(runtimeSessionId: string, maxBytes: number): Promise { + const key = await this.latestKey(runtimeSessionId); + if (!key) return null; + let size: number; + try { + const stat = await this.client.statObject(this.bucket, key); + size = stat.size; + } catch (error) { + const code = (error as { code?: string })?.code; + if (code === 'NoSuchKey' || code === 'NotFound') return null; + throw error; + } + /* Reject before downloading — never buffer an arbitrarily large object. */ + if (size > maxBytes) { + throw new CheckpointTooLargeError(`checkpoint ${size}B exceeds maxBytes ${maxBytes}B`); + } + const stream = await this.client.getObject(this.bucket, key); + const chunks: Buffer[] = []; + let total = 0; + for await (const chunk of stream) { + /* Cap during download too: the object can grow between statObject and here + * (a concurrent put), and the stat size alone wouldn't catch it. */ + total += (chunk as Buffer).length; + if (total > maxBytes) { + throw new CheckpointTooLargeError(`checkpoint exceeded maxBytes ${maxBytes}B during download`); + } + chunks.push(chunk as Buffer); + } + return Buffer.concat(chunks); + } + + private async listKeys(runtimeSessionId: string): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + const keys: string[] = []; + const stream = this.client.listObjectsV2(this.bucket, prefix, true); + for await (const item of stream as AsyncIterable<{ name?: string }>) { + if (item.name) keys.push(item.name); + } + return keys; + } + + private async latestKey(runtimeSessionId: string): Promise { + const keys = await this.listKeys(runtimeSessionId); + if (keys.length === 0) return null; + return keys.reduce((max, key) => (key > max ? key : max)); + } + + private async pruneOlderThan(runtimeSessionId: string, keepKey: string): Promise { + const stale = (await this.listKeys(runtimeSessionId)).filter((key) => key < keepKey); + if (stale.length > 0) await this.client.removeObjects(this.bucket, stale); + } +} + +/** 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(); + + async put(runtimeSessionId: string, sequence: number, data: Buffer): Promise { + const key = checkpointObjectKey(runtimeSessionId, sequence); + const prefix = checkpointPrefixFor(runtimeSessionId); + for (const existing of this.objects.keys()) { + /* prune strictly-older checkpoints only — never a newer one */ + if (existing.startsWith(prefix) && existing < key) this.objects.delete(existing); + } + this.objects.set(key, Buffer.from(data)); + } + + 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): Promise { + const prefix = checkpointPrefixFor(runtimeSessionId); + let latest: string | undefined; + for (const key of this.objects.keys()) { + if (key.startsWith(prefix) && (latest === undefined || key > latest)) latest = key; + } + if (latest === undefined) return null; + const data = this.objects.get(latest) as Buffer; + if (data.length > maxBytes) { + throw new CheckpointTooLargeError(`checkpoint ${data.length}B exceeds maxBytes ${maxBytes}B`); + } + return 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..afb2351 --- /dev/null +++ b/service/src/runtime-session/checkpoint.ts @@ -0,0 +1,231 @@ +import axios from 'axios'; +import type { MicrovmAuthToken } from './lambda-client'; +import { microvmPortHeaders } from './lambda-client'; +import type { CheckpointStore } from './checkpoint-store'; +import { + acquireRuntimeSessionLock, + allocateCheckpointSequence, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + reseedCheckpointSequence, + writeRuntimeSessionRecord, +} from './registry'; +import { checkpointObjectKey } from './checkpoint-store'; +import { microvmCheckpoints, microvmRestores, microvmCheckpointBytes } from '../metrics'; +import logger from '../logger'; + +/** Reject if `promise` doesn't settle within `ms`. The underlying op is not + * cancelled (the object-store client has no abort hook), but the caller stops + * waiting so a stalled write can't hold the session lock indefinitely. */ +function withTimeout(promise: Promise, ms: number, label: string): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error(`${label} timed out after ${ms}ms`)), ms); + promise.then( + (value) => { clearTimeout(timer); 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. Failures are never fatal: a missed checkpoint + * degrades to file-ref recovery, a failed restore degrades to a fresh + * workspace ("a resumed VM can be faster, but a relaunched VM must be + * correct"). + */ + +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 }, + 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: 'arraybuffer', + maxContentLength: config.maxBytes, + timeout: config.timeoutMs, + }); + return Buffer.from(response.data); +} + +export async function pushRestore( + args: { mintToken: () => Promise; endpointBase: string; runtimeSessionId: string }, + data: Buffer, + config: CheckpointConfig, +): Promise { + const token = await args.mintToken(); + await axios.post(`${args.endpointBase}/api/v2/session/restore`, data, { + headers: { + [token.headerName]: token.token, + ...microvmPortHeaders(config.port), + [RUNTIME_SESSION_ID_HEADER]: args.runtimeSessionId, + 'Content-Type': 'application/x-gtar', + }, + maxBodyLength: config.maxBytes, + timeout: config.timeoutMs, + }); +} + +/** + * 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; +}): Promise<'stored' | 'skipped_busy' | 'skipped_state' | 'failed'> { + const heldToken = args.lockToken; + const lockToken = heldToken ?? await acquireRuntimeSessionLock(args.runtimeSessionId); + 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; + const data = await pullCheckpoint({ + mintToken: () => args.mintToken(microvmId), + endpointBase: args.normalizeEndpoint(record.endpoint), + runtimeSessionId: args.runtimeSessionId, + }, args.config); + /* Allocate the checkpoint's sequence. INCR==1 means a fresh OR TTL-reset + * counter, so seed it above any objects still retained in the store (else a + * post-idle checkpoint would write seq 1 and restore would keep picking a + * stale higher-sequence object). Checkpoints for one session serialize on + * the lock, so this read-then-seed has no concurrent writer. */ + let sequence = await allocateCheckpointSequence(args.runtimeSessionId); + if (sequence === 1) { + const retainedMax = await args.store.latestSequence(args.runtimeSessionId); + if (retainedMax >= sequence) { + sequence = retainedMax + 1; + await reseedCheckpointSequence(args.runtimeSessionId, sequence); + } + } + /* Fence the pointer write first: if it reports we were fenced, skip the + * store entirely. */ + 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'; + } + /* Bound the object-store write by the checkpoint timeout too — otherwise a + * stalled S3/MinIO put holds the session lock past JOB_TIMEOUT. */ + await withTimeout( + args.store.put(args.runtimeSessionId, sequence, data), + args.config.timeoutMs, + 'checkpoint store.put', + ); + microvmCheckpointBytes.observe(data.length); + 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 { + /* 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; +}): Promise<'restored' | 'absent' | 'fetch_failed' | 'push_failed'> { + /* The store enforces `maxBytes` before buffering (stats S3 object size first), + * so an oversized/stray checkpoint throws here instead of OOM'ing the worker. + * 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: Buffer | null; + try { + data = await withTimeout( + args.store.get(args.runtimeSessionId, args.config.maxBytes), + args.config.timeoutMs, + 'checkpoint store.get', + ); + } catch (error) { + /* Fetch failed before the runner was touched — the workspace is still the + * clean fresh-VM one, so the caller can safely execute (degraded, no + * restore). */ + microvmRestores.inc({ outcome: 'failed' }); + logger.warn('Checkpoint fetch failed; continuing with a fresh 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, + }, data, args.config); + microvmRestores.inc({ outcome: 'restored' }); + logger.info('Session workspace restored from checkpoint', { + runtimeSessionId: args.runtimeSessionId, + bytes: data.length, + }); + 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'; + } +} diff --git a/service/src/runtime-session/id.test.ts b/service/src/runtime-session/id.test.ts new file mode 100644 index 0000000..18f3412 --- /dev/null +++ b/service/src/runtime-session/id.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, test } from 'bun:test'; +import { + RUNTIME_SESSION_HINT_MAX_LENGTH, + RuntimeSessionHintError, + deriveRuntimeSessionId, + 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); + }); +}); diff --git a/service/src/runtime-session/id.ts b/service/src/runtime-session/id.ts new file mode 100644 index 0000000..eb2a441 --- /dev/null +++ b/service/src/runtime-session/id.ts @@ -0,0 +1,59 @@ +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. */ +export function resolveRuntimeSessionIdForRequest(args: { + mode: 'stateless' | 'affinity' | 'strict'; + storageNamespace: string; + canonicalUserId: string; + hint?: string; +}): string | undefined { + if (args.mode === 'stateless') return undefined; + return deriveRuntimeSessionId(args); +} 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..3031f9f --- /dev/null +++ b/service/src/runtime-session/lambda-client-aws.test.ts @@ -0,0 +1,173 @@ +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', + '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}`); + }); +}); + +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..0ee48b6 --- /dev/null +++ b/service/src/runtime-session/lambda-client-aws.ts @@ -0,0 +1,171 @@ +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): Promise; +} + +export class AwsLambdaMicrovmClient implements LambdaMicrovmClient { + private readonly client: MicrovmCommandSender; + + constructor(options: { region?: string; client?: MicrovmCommandSender } = {}) { + this.client = options.client ?? new LambdaMicrovmsClient({ + region: options.region, + retryMode: 'adaptive', + maxAttempts: 3, + }); + } + + private async send(operation: string, command: unknown): Promise { + try { + return await this.client.send(command) as T; + } catch (error) { + throw new LambdaMicrovmApiError( + classifyError(error), + operation, + (error as Error)?.message ?? `Lambda MicroVM ${operation} failed`, + error, + ); + } + } + + async runMicrovm(args: RunMicrovmArgs): Promise { + const response = await this.send[0]>('RunMicrovm', new RunMicrovmCommand({ + 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, + })); + return toDescription(response); + } + + async getMicrovm(microvmId: string): Promise { + const response = await this.send[0]>( + 'GetMicrovm', + new GetMicrovmCommand({ microvmIdentifier: microvmId }), + ); + return toDescription(response); + } + + async suspendMicrovm(microvmId: string): Promise { + await this.send('SuspendMicrovm', new SuspendMicrovmCommand({ microvmIdentifier: microvmId })); + } + + async resumeMicrovm(microvmId: string): Promise { + const response = await this.send[0]>( + 'ResumeMicrovm', + new ResumeMicrovmCommand({ microvmIdentifier: microvmId }), + ); + return toDescription(response); + } + + async terminateMicrovm(microvmId: string): Promise { + await this.send('TerminateMicrovm', new TerminateMicrovmCommand({ microvmIdentifier: microvmId })); + } + + async createMicrovmAuthToken(args: { + microvmId: string; + port: number; + ttlSeconds: number; + }): 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 }], + }), + ); + 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..bc96b36 --- /dev/null +++ b/service/src/runtime-session/lambda-client-fake.ts @@ -0,0 +1,192 @@ +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); + } + + 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 vm: FakeMicrovm = { + microvmId, + state: 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..676c7e1 --- /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): Promise; + getMicrovm(microvmId: string): Promise; + suspendMicrovm(microvmId: string): Promise; + resumeMicrovm(microvmId: string): Promise; + terminateMicrovm(microvmId: string): Promise; + createMicrovmAuthToken(args: { + microvmId: string; + port: number; + ttlSeconds: number; + }): Promise; +} diff --git a/service/src/runtime-session/registry.test.ts b/service/src/runtime-session/registry.test.ts new file mode 100644 index 0000000..cf30f26 --- /dev/null +++ b/service/src/runtime-session/registry.test.ts @@ -0,0 +1,155 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import { + acquireRuntimeSessionLock, + allocateRuntimeSessionGeneration, + countActiveRuntimeSessions, + forgetRuntimeSessionActive, + listIdleRuntimeSessions, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + removeRuntimeSession, + resetRedisForTests, + setRedisForTests, + touchRuntimeSessionActive, + 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); + }); +}); + +describe('fenced record writes', () => { + 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 record + active member', async () => { + const token = (await acquireRuntimeSessionLock('rt_abc123')) as string; + await writeRuntimeSessionRecord(record(), token); + await touchRuntimeSessionActive('rt_abc123', 1_778_250_000_000); + + 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(); + expect(await countActiveRuntimeSessions()).toBe(0); + }); +}); + +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('active session bookkeeping', () => { + test('idle listing returns only sessions last seen before the cutoff', async () => { + await touchRuntimeSessionActive('rt_old', 1_000); + await touchRuntimeSessionActive('rt_mid', 5_000); + await touchRuntimeSessionActive('rt_new', 9_000); + + expect(await listIdleRuntimeSessions(4_999)).toEqual(['rt_old']); + expect(await listIdleRuntimeSessions(5_000)).toEqual(['rt_old', 'rt_mid']); + expect(await countActiveRuntimeSessions()).toBe(3); + }); + + test('touch updates the score in place; forget repairs orphans', async () => { + await touchRuntimeSessionActive('rt_abc123', 1_000); + await touchRuntimeSessionActive('rt_abc123', 9_000); + expect(await listIdleRuntimeSessions(5_000)).toEqual([]); + expect(await countActiveRuntimeSessions()).toBe(1); + + await forgetRuntimeSessionActive('rt_abc123'); + expect(await countActiveRuntimeSessions()).toBe(0); + }); + + test('idle listing respects the limit bound', async () => { + for (let i = 0; i < 5; i++) { + await touchRuntimeSessionActive(`rt_${i}`, i); + } + expect(await listIdleRuntimeSessions(10, 2)).toHaveLength(2); + }); +}); diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts new file mode 100644 index 0000000..578f084 --- /dev/null +++ b/service/src/runtime-session/registry.ts @@ -0,0 +1,283 @@ +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) + * rtsx:active zset of session ids by last_seen_at (sweeper-pruned) + * + * 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 the ingress/egress network connector ARNs the VM launched + * with — connectors are only applied at RunMicrovm, so a config change must + * make an existing session non-reusable (else a tightened egress policy is + * bypassed by warm reuse). */ + connectors?: 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:'; +const ACTIVE_ZSET = 'rtsx:active'; + +/** 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, + activeKey: string, + token: string, + member: 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: 3, + lua: `if redis.call('get', KEYS[2]) == ARGV[1] then + redis.call('del', KEYS[1]) + redis.call('zrem', KEYS[3], ARGV[2]) + return 1 +else + return 0 +end`, + }); + 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. + * Callers decide mode policy (affinity ⇒ stateless fallback, strict ⇒ 409). */ +export async function waitForRuntimeSessionLock( + runtimeSessionId: string, + args: { waitMs: number; pollMs?: number; ttlMs?: number }, +): Promise { + const pollMs = args.pollMs ?? 250; + const deadline = Date.now() + args.waitMs; + for (;;) { + const token = await acquireRuntimeSessionLock(runtimeSessionId, args.ttlMs); + if (token != null) return token; + if (Date.now() + pollMs > deadline) return null; + await new Promise((resolve) => setTimeout(resolve, pollMs)); + } +} + +export async function releaseRuntimeSessionLock(runtimeSessionId: string, token: string): Promise { + try { + await redis.releaseRuntimeSessionLockScript(`${LOCK_PREFIX}${runtimeSessionId}`, token); + } catch (err) { + logger.warn('Failed to release runtime session lock', { runtimeSessionId, err }); + } +} + +/** 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. */ +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; + } catch (err) { + logger.warn('Failed to renew runtime session lock', { runtimeSessionId, err }); + return false; + } +} + +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; +} + +/** Monotonic per-checkpoint sequence via INCR. Bounded by the record TTL, so it + * can reset after a long idle; the caller re-seeds it above any retained + * objects (see reseedCheckpointSequence). A pure counter — no wall clock — so + * ordering is unaffected by cross-pod clock skew. */ +export async function allocateCheckpointSequence(runtimeSessionId: string): Promise { + const key = `${CKPT_SEQ_PREFIX}${runtimeSessionId}`; + const sequence = await redis.incr(key); + await redis.expire(key, RUNTIME_SESSION_RECORD_TTL_SECONDS); + return sequence; +} + +/** Force the checkpoint counter up to `value` after a TTL reset dropped it below + * the sequences still retained in the object store, so the next INCR continues + * above them and restore never picks a stale higher-sequence object. */ +export async function reseedCheckpointSequence(runtimeSessionId: string, value: number): Promise { + const key = `${CKPT_SEQ_PREFIX}${runtimeSessionId}`; + await redis.set(key, String(value), 'EX', RUNTIME_SESSION_RECORD_TTL_SECONDS); +} + +export async function touchRuntimeSessionActive(runtimeSessionId: string, lastSeenAtMs: number): Promise { + await redis.zadd(ACTIVE_ZSET, lastSeenAtMs, runtimeSessionId); +} + +/** Fenced removal: deletes the record and active-zset member 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}`, + ACTIVE_ZSET, + lockToken, + runtimeSessionId, + ); + return result === 1; +} + +/** Unfenced zset repair for sweeper use (record already gone). */ +export async function forgetRuntimeSessionActive(runtimeSessionId: string): Promise { + await redis.zrem(ACTIVE_ZSET, runtimeSessionId); +} + +export async function listIdleRuntimeSessions(idleBeforeMs: number, limit = 100): Promise { + return redis.zrangebyscore(ACTIVE_ZSET, '-inf', idleBeforeMs, 'LIMIT', 0, limit); +} + +export async function countActiveRuntimeSessions(): Promise { + return redis.zcard(ACTIVE_ZSET); +} diff --git a/service/src/runtime-session/throttle.test.ts b/service/src/runtime-session/throttle.test.ts new file mode 100644 index 0000000..3c26cf8 --- /dev/null +++ b/service/src/runtime-session/throttle.test.ts @@ -0,0 +1,89 @@ +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 }); + }); +}); diff --git a/service/src/runtime-session/throttle.ts b/service/src/runtime-session/throttle.ts new file mode 100644 index 0000000..2694c2d --- /dev/null +++ b/service/src/runtime-session/throttle.ts @@ -0,0 +1,87 @@ +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) => Promise; +} + +const defaultSleep = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)); + +/** + * 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 (;;) { + const poisoned = await redis.pttl(`${POISON_PREFIX}${op}`); + if (poisoned > 0) { + if (now() + poisoned > deadline) throw new MicrovmOpThrottledError(op, options.budgetMs); + await sleep(poisoned); + continue; + } + + const nowMs = now(); + const second = Math.floor(nowMs / 1_000); + const key = `${BUCKET_PREFIX}${op}:${second}`; + const count = await redis.incr(key); + if (count === 1) { + await redis.pexpire(key, BUCKET_TTL_MS); + } + 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); + } +} + +/** 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..3061348 --- /dev/null +++ b/service/src/sandbox-backend/index.test.ts @@ -0,0 +1,40 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +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 lambda-microvm backend when configured', async () => { + env.SANDBOX_BACKEND = 'lambda-microvm'; + const backend = getSandboxBackend(); + const { LambdaMicrovmSandboxBackend } = await import('./lambda-microvm'); + expect(backend).toBeInstanceOf(LambdaMicrovmSandboxBackend); + expect(backend.name).toBe('lambda-microvm'); + }); + + 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..481ee4c --- /dev/null +++ b/service/src/sandbox-backend/index.ts @@ -0,0 +1,63 @@ +import type { SandboxBackend } from './types'; +import { LambdaMicrovmSandboxBackend } from './lambda-microvm'; +import { HttpSandboxBackend } from './http'; +import { MinioCheckpointStore } from '../runtime-session/checkpoint-store'; +import { env } from '../config'; + +export type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +export { SandboxBackendError } from './types'; +export { HttpSandboxBackend } from './http'; +export { LambdaMicrovmSandboxBackend } from './lambda-microvm'; + +let backend: SandboxBackend | undefined; + +function createBackend(): SandboxBackend { + if (env.SANDBOX_BACKEND === 'lambda-microvm') { + return new LambdaMicrovmSandboxBackend({ + /* Dynamic import keeps @aws-sdk out of http-only worker bundles. */ + clientFactory: async () => { + 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: env.SESSION_CHECKPOINTS && env.RUNTIME_SESSION_MODE !== 'stateless' + ? new MinioCheckpointStore() + : undefined, + }); + } + 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..9b2888a --- /dev/null +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -0,0 +1,587 @@ +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import axios from 'axios'; +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, + resetRedisForTests as resetRegistryRedis, + setRedisForTests as setRegistryRedis, + writeRuntimeSessionRecord, +} from '../runtime-session/registry'; +import { MemoryCheckpointStore, checkpointObjectKey, checkpointPrefixFor } from '../runtime-session/checkpoint-store'; +import { LambdaMicrovmSandboxBackend, normalizeMicrovmEndpoint, type LambdaMicrovmBackendConfig } from './lambda-microvm'; +import { SandboxBackendError } from './types'; +import type { SandboxExecuteContext, 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 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; + captured.push({ + path, + rawBody: await req.text(), + headers: Object.fromEntries(req.headers.entries()), + }); + if (path === '/api/v2/health') { + return new Response(JSON.stringify({ status: 'ok' }), { + status: healthStatus, + headers: { 'Content-Type': 'application/json' }, + }); + } + if (path === '/api/v2/execute') { + if (executeDelayMs > 0) { + await new Promise((resolve) => setTimeout(resolve, executeDelayMs)); + } + return new Response(JSON.stringify(EXECUTE_RESPONSE), { + 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') { + return new Response(JSON.stringify({ status: 'restored' }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response('not found', { status: 404 }); + }, + }); +}); + +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; +}); + +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]; + expect(executeReq?.headers['x-aws-proxy-auth']).toBe(vm.mintedTokens[0]); + + 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('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('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); + }); +}); + +describe('LambdaMicrovmSandboxBackend session execution', () => { + 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('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 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); + }); + + 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('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('affinity mode falls back to a stateless one-shot when the lock is held', async () => { + const fake = fakeClient(); + const backend = makeBackend(fake, { lockWaitMs: 100 }); + await acquireRuntimeSessionLock('rt_session_1', 60_000); + + const result = await backend.execute(request(), sessionContext({ runtimeSessionMode: 'affinity' })); + expect(result).toEqual(EXECUTE_RESPONSE); + /* Stateless fallback: launched a one-shot VM and terminated it. */ + const runArgs = fake.callsFor('runMicrovm')[0].args as { runHookPayload?: string; clientToken: string }; + expect(runArgs.runHookPayload).toBeUndefined(); + expect(runArgs.clientToken.startsWith('exec-')).toBe(true); + expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); + }); + + 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('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('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'); + expect((await store.get('rt_ckpt_1', 1_000_000))?.toString()).toBe(checkpointBlob); + const record = await readRuntimeSessionRecord('rt_ckpt_1'); + /* Key is timestamp-based now (Date.now at checkpoint), so match the shape. */ + expect(record?.workspace_checkpoint).toStartWith(checkpointPrefixFor('rt_ckpt_1')); + expect(record?.workspace_checkpoint).toEndWith('.tar.gz'); + expect(record?.checkpointed_at).toBeGreaterThan(0); + }); + + 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')); + /* 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(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()); + expect(captured.filter((c) => c.path.startsWith('/api/v2/session/'))).toHaveLength(0); + expect(store.objects.size).toBe(0); + }); + + 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(); + }); +}); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts new file mode 100644 index 0000000..a10a9bc --- /dev/null +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -0,0 +1,727 @@ +import axios from 'axios'; +import { nanoid } from 'nanoid'; +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, restoreSession } from '../runtime-session/checkpoint'; +import { + RUNTIME_SESSION_LOCK_TTL_MS, + allocateRuntimeSessionGeneration, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + removeRuntimeSession, + renewRuntimeSessionLock, + touchRuntimeSessionActive, + waitForRuntimeSessionLock, + writeRuntimeSessionRecord, +} from '../runtime-session/registry'; +import { + microvmLaunches, + microvmLaunchDuration, + microvmTerminations, + microvmThrottleEvents, + runtimeSessionFallback, + runtimeSessionLockContention, +} from '../metrics'; +import { injectTraceHeaders, withSpan } from '../telemetry'; +import { SandboxBackendError } from './types'; +import { Jobs } from '../enum'; +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; +} + +interface LambdaMicrovmBackendDeps { + clientFactory: () => Promise; + config: LambdaMicrovmBackendConfig; + pollIntervalMs?: number; + /** Injected in session+checkpoint mode; undefined disables checkpoints. */ + checkpointStore?: CheckpointStore; +} + +const sleep = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)); + +/** 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 { + const client = await this.client(); + if (ctx.runtimeSessionId && ctx.runtimeSessionMode !== 'stateless') { + return this.executeSession(client, req, ctx, ctx.runtimeSessionId); + } + 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 { + return await this.proxyExecute(client, vm, req, ctx); + } 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, + ): Promise { + /* Approximate the JOB_TIMEOUT budget consumed so we don't push the result + * past the router's `waitUntilFinished(JOB_TIMEOUT)` with an optional + * checkpoint (below). Captured before the lock wait so lock-wait + launch + + * execute all count. */ + const startedAt = Date.now(); + const lockToken = await waitForRuntimeSessionLock(runtimeSessionId, { waitMs: this.config.lockWaitMs }); + if (!lockToken) { + runtimeSessionLockContention.inc({ mode: ctx.runtimeSessionMode }); + if (ctx.runtimeSessionMode === 'strict') { + throw new SandboxBackendError('RUNTIME_SESSION_BUSY', `Runtime session ${runtimeSessionId} is busy`); + } + /* Affinity: warmth is only an optimization — fall back to a correct + * stateless one-shot (the payload still carries all file refs). */ + runtimeSessionFallback.inc(); + return this.executeStateless(client, req, ctx); + } + + /* Heartbeat the lock for as long as we hold it so an arbitrarily long + * critical path (launch throttle + readiness/restore + execute + checkpoint, + * each with its own I/O and token-mint waits) can't outlive a fixed TTL and + * let a second worker fence us and run concurrently. Fenced renew stops + * itself; the interval is a third of the TTL so a couple of missed ticks are + * survivable. */ + const heartbeat = setInterval(() => { + void renewRuntimeSessionLock(runtimeSessionId, lockToken).then((held) => { + if (!held) clearInterval(heartbeat); + }); + }, Math.floor(RUNTIME_SESSION_LOCK_TTL_MS / 3)); + + try { + const existing = await readRuntimeSessionRecord(runtimeSessionId); + const { vm, reused } = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); + const result = await this.executeOnSessionVm(client, vm, req, ctx, runtimeSessionId, lockToken, reused); + /* 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 = this.config.jobTimeoutMs - (now - startedAt); + /* Reserve the WHOLE checkpoint path, not just one timeout: it can spend a + * token-budget wait (up to launchTimeoutMs) + the checkpoint GET + the + * object-store put (each up to checkpoint.timeoutMs). Guarding on a single + * timeout let a run finishing with barely more than that still block long + * enough to blow the router's waitUntilFinished(JOB_TIMEOUT) after the + * sandbox work already succeeded. */ + const worstCaseCheckpointMs = + this.config.launchTimeoutMs + 2 * this.config.checkpoint.timeoutMs; + const canCheckpoint = !ctx.signal.aborted && remainingBudgetMs > worstCaseCheckpointMs; + const settled = await readRuntimeSessionRecord(runtimeSessionId); + const nextRecord = settled + ? canCheckpoint + ? await this.checkpointUnderLock(client, settled, runtimeSessionId, now, lockToken) + : { ...settled, state: 'RUNNING' as const, last_seen_at: now } + : undefined; + if (nextRecord) { + await writeRuntimeSessionRecord(nextRecord, lockToken); + } + await touchRuntimeSessionActive(runtimeSessionId, now); + return result; + } finally { + clearInterval(heartbeat); + 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 drops the registry record 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, + reused: boolean, + ): Promise { + try { + return await this.proxyExecute(client, vm, req, ctx, runtimeSessionId, reused); + } 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) on a REUSED VM is the AWS gateway + * reporting the VM as unreachable — typically a suspended VM that failed to + * auto-resume. 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; + const transportFailure = axios.isAxiosError(error) && error.response == null; + const unhealthy = error instanceof SandboxBackendError && error.code === 'MICROVM_UNHEALTHY'; + const gatewayUnreachable = reused && status >= 502 && status <= 504; + if (ctx.signal.aborted || transportFailure || unhealthy || gatewayUnreachable) { + await this.terminate(client, vm.microvmId, ctx.signal.aborted ? 'timeout' : 'error').catch(() => {}); + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + } + throw error; + } + } + + /** Order-independent fingerprint of the ingress/egress connector ARNs the + * current config would launch a VM with, recorded on the session so a + * connector config change makes an existing VM non-reusable. */ + private connectorFingerprint(): string { + const ingress = [...(this.config.ingressConnectorArns ?? [])].sort(); + const egress = [...(this.config.egressConnectorArns ?? [])].sort(); + return JSON.stringify({ ingress, egress }); + } + + private async findOrLaunchSession( + client: LambdaMicrovmClient, + ctx: SandboxExecuteContext, + runtimeSessionId: string, + record: RuntimeSessionRecord | null, + lockToken: string, + ): Promise<{ vm: MicrovmDescription; reused: boolean }> { + 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.connectors ?? '') === this.connectorFingerprint(); + /* 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 }, + reused: true, + }; + } + + /* We're relaunching. If the recorded VM is a live-but-non-reusable one + * (config/version/port drift, or deadline too close for this job) it would + * otherwise leak — running/suspended and billing until idle/max-duration + * expiry — once we overwrite the record below. Terminate it first. Skip + * when likelyIdleTerminated: AWS already killed it, so that's positive + * evidence it's gone and terminate would be a wasted not-found call. */ + if (record?.microvm_id && !likelyIdleTerminated) { + await this.terminate(client, record.microvm_id, 'superseded').catch(() => {}); + } + + const generation = await allocateRuntimeSessionGeneration(runtimeSessionId); + const pendingOk = await writeRuntimeSessionRecord({ + runtime_session_id: runtimeSessionId, + tenant_id: ctx.tenantId ?? '', + canonical_user_id: ctx.canonicalUserId ?? '', + state: 'PENDING', + generation, + last_seen_at: Date.now(), + }, lockToken); + if (!pendingOk) { + throw new SandboxBackendError('MICROVM_FENCED', `Lost session lock for ${runtimeSessionId} before launch`); + } + + const launchedAt = Date.now(); + 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 runningOk = await writeRuntimeSessionRecord({ + runtime_session_id: runtimeSessionId, + tenant_id: ctx.tenantId ?? '', + canonical_user_id: ctx.canonicalUserId ?? '', + microvm_id: vm.microvmId, + endpoint: vm.endpoint, + port: this.config.port, + image_arn: this.config.imageArn, + image_version: this.config.imageVersion, + connectors: this.connectorFingerprint(), + state: 'RUNNING', + generation, + launched_at: launchedAt, + last_seen_at: Date.now(), + hard_deadline_at: launchedAt + this.config.maxDurationSeconds * 1_000 - 60_000, + }, lockToken); + if (!runningOk) { + await this.terminate(client, vm.microvmId, 'error'); + throw new SandboxBackendError('MICROVM_FENCED', `Lost session lock for ${runtimeSessionId} after launch`); + } + + /* Fresh VM: restore the predecessor's workspace before the first execute so + * an 8h rollover / eviction is invisible. Attempt whenever checkpoints are + * active, NOT only when a Redis record exists — the checkpoint lives under a + * deterministic S3 key that can outlive/repair a lost record, and + * restoreSession treats a missing object as `absent` (a truly new session + * just no-ops one stat). */ + if (this.checkpointStore && this.checkpointsActive()) { + const endpointBase = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + /* Wait for the runner's API listener before restoring. RUNNING is a + * control-plane state — the endpoint can be allocated while the app is + * still booting, and pushRestore is intentionally non-fatal, so a restore + * that raced the boot would silently drop the checkpoint and run the first + * execute on an empty workspace. If the runner never comes up, tear the VM + * down so the next call relaunches instead of reusing a dead endpoint. */ + try { + await this.waitForRunnerReady(client, vm.microvmId, endpointBase, ctx); + } catch (error) { + await this.terminate(client, vm.microvmId, 'error').catch(() => {}); + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + throw error; + } + const restoreResult = await restoreSession({ + mintToken: (microvmId) => this.mintAuthToken(client, microvmId), + store: this.checkpointStore, + runtimeSessionId, + microvmId: vm.microvmId, + endpointBase, + config: this.config.checkpoint, + }); + /* A push-restore that failed after the runner began extracting leaves the + * workspace possibly-partial (the runner's cleanup runs async past our + * abort). Don't execute against it — recycle so the next call relaunches a + * clean VM. (A fetch failure / absent checkpoint is safe: the workspace is + * the untouched fresh one.) */ + if (restoreResult === 'push_failed') { + await this.terminate(client, vm.microvmId, 'error').catch(() => {}); + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + throw new SandboxBackendError('MICROVM_UNHEALTHY', 'Checkpoint restore left the workspace in an unknown state'); + } + } + return { vm, reused: false }; + } + + /** 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 token = await this.mintAuthToken(client, microvmId); + const deadline = Date.now() + this.config.launchTimeoutMs; + let lastError: unknown; + while (Date.now() < deadline) { + if (ctx.signal.aborted) { + throw new SandboxBackendError('MICROVM_LAUNCH_FAILED', 'Execution aborted while waiting for runner readiness'); + } + try { + await this.assertHealthy(base, token.token, ctx); + return; + } catch (error) { + lastError = error; + await sleep(this.pollIntervalMs); + } + } + 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 checkpoint degrades to file-ref recovery. + */ + private async checkpointUnderLock( + client: LambdaMicrovmClient, + record: RuntimeSessionRecord, + runtimeSessionId: string, + now: number, + lockToken: string, + ): 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), + store: this.checkpointStore, + runtimeSessionId, + config: this.config.checkpoint, + normalizeEndpoint: normalizeMicrovmEndpoint, + lockToken, + }); + /* 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): Promise { + try { + await acquireOpBudget('token', { + limitPerSecond: this.config.tokenTps, + budgetMs: this.config.launchTimeoutMs, + }); + } 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, + }); + } 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; + } + } + + private async proxyExecute( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + runtimeSessionId?: string, + reused = false, + ): Promise { + const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + const token = await this.mintAuthToken(client, vm.microvmId); + /* 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 (!reused) { + await this.assertHealthy(base, token.token, 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 { + const endLaunchTimer = microvmLaunchDuration.startTimer(); + try { + await acquireOpBudget('run', { + limitPerSecond: this.config.launchTps, + budgetMs: this.config.launchTimeoutMs, + }); + } 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, + }); + } 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' }); + await this.terminate(client, vm.microvmId, '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') { + throw new SandboxBackendError( + 'MICROVM_LAUNCH_FAILED', + `MicroVM ${current.microvmId} entered ${current.state} before becoming ready`, + ); + } + 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); + current = await client.getMicrovm(current.microvmId); + } + 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); + microvmTerminations.inc({ reason }); + } catch (error) { + logger.error('Failed to terminate MicroVM', { microvmId, reason, error }); + } + } +} diff --git a/service/src/sandbox-backend/types.ts b/service/src/sandbox-backend/types.ts new file mode 100644 index 0000000..190ab70 --- /dev/null +++ b/service/src/sandbox-backend/types.ts @@ -0,0 +1,59 @@ +import type * as t from '../types'; + +/** + * 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; +} + +export interface SandboxExecuteContext { + executionId: string; + language: string; + isSynthetic: boolean; + /** Worker-owned JOB_TIMEOUT abort signal. */ + signal: AbortSignal; + tenantId?: string; + canonicalUserId?: string; + /** Absent ⇒ stateless execution (no runtime session affinity). */ + runtimeSessionId?: string; + runtimeSessionMode: 'stateless' | 'affinity' | 'strict'; +} + +/** 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 { + constructor( + public readonly code: SandboxBackendErrorCode, + message: string, + public readonly cause?: unknown, + ) { + super(message); + this.name = 'SandboxBackendError'; + } +} diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 15224ad..705c8ae 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -3,12 +3,20 @@ 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, + lambdaEgressConnectors: env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS, + lambdaTokenTtl: env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS, + lambdaAllowShell: env.LAMBDA_MICROVM_ALLOW_SHELL, gatewayUrl: env.EGRESS_GATEWAY_URL, grantSecret: env.EGRESS_GRANT_SECRET, privateKey: env.EXECUTION_MANIFEST_PRIVATE_KEY, @@ -24,6 +32,13 @@ 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_EGRESS_CONNECTOR_ARNS = saved.lambdaEgressConnectors; + env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS = saved.lambdaTokenTtl; + env.LAMBDA_MICROVM_ALLOW_SHELL = saved.lambdaAllowShell; env.EGRESS_GATEWAY_URL = saved.gatewayUrl; env.EGRESS_GRANT_SECRET = saved.grantSecret; env.EXECUTION_MANIFEST_PRIVATE_KEY = saved.privateKey; @@ -120,3 +135,89 @@ 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_EGRESS_CONNECTOR_ARNS = undefined; + env.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS = 300; + env.LAMBDA_MICROVM_ALLOW_SHELL = false; + /* 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('rejects session checkpoints without object storage configured', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + delete process.env.MINIO_ACCESS_KEY; + 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('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('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 900 or less'); + + 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..8093d77 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -1,4 +1,5 @@ import { env } from './config'; +import logger from './logger'; import { INTERNAL_SERVICE_TOKEN_ENV } from './internal-service-auth'; export class SecureStartupConfigError extends Error { @@ -48,6 +49,69 @@ 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; + + 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.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.LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS > 900) { + throw new SecureStartupConfigError('LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS must be 900 or less'); + } + 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/empty creds, 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). */ + if (env.SESSION_CHECKPOINTS && env.RUNTIME_SESSION_MODE !== 'stateless') { + const missing = ['MINIO_ENDPOINT', 'MINIO_ACCESS_KEY', 'MINIO_SECRET_KEY'].filter( + (name) => !nonEmpty(process.env[name]), + ); + 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/router.ts b/service/src/service/router.ts index 9d9fd37..f914662 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 { resolveRuntimeSessionIdForRequest, validateRuntimeSessionHint, RuntimeSessionHintError } from '../runtime-session/id'; import { jobsSubmitted } from '../metrics'; import { captureTraceCarrier, withSpan } from '../telemetry'; import { Jobs, Languages } from '../enum'; @@ -133,6 +134,27 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) return res.status(400).json({ error: `Unsupported language: ${rawLang}` }); } + let runtimeSessionId: string | undefined; + try { + /* In stateless mode the hint is ignored entirely, so don't validate it — + * a malformed hint that will never be used must not 400 the request. */ + const hint = + env.RUNTIME_SESSION_MODE === 'stateless' + ? undefined + : validateRuntimeSessionHint(body.runtime_session_hint); + runtimeSessionId = resolveRuntimeSessionIdForRequest({ + mode: env.RUNTIME_SESSION_MODE, + storageNamespace: identity.storageNamespace, + canonicalUserId: identity.canonicalUserId, + hint, + }); + } 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 +241,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..b8a0ec9 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 }; @@ -227,6 +234,8 @@ export type JobData = { executionId?: string; tenantId?: string; canonicalUserId?: string; + /** Server-derived runtime session identity (absent ⇒ stateless execution). */ + runtimeSessionId?: string; 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..ea5ccf9 100644 --- a/service/src/utils.test.ts +++ b/service/src/utils.test.ts @@ -129,6 +129,22 @@ 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 rt_abc 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: 'MicroVM did not reach RUNNING within 60000ms' }, + }); + }); + 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..156c651 100644 --- a/service/src/utils.ts +++ b/service/src/utils.ts @@ -108,6 +108,17 @@ 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 : ''; + + /* Lambda MicroVM backend failures are surfaced by the worker as + * `: `. A busy runtime session in strict mode is a 409; + * launch/health failures are transient upstream errors (503). */ + const backendMatch = message.match(/^(RUNTIME_SESSION_BUSY|MICROVM_[A-Z_]+):\s*(.+)$/); + if (backendMatch) { + const code = backendMatch[1]; + const status = code === 'RUNTIME_SESSION_BUSY' ? 409 : 503; + return { status, body: { error: code.toLowerCase(), message: backendMatch[2] } }; + } + 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..f15c393 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -3,26 +3,21 @@ 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 { summarizeSandboxResponse, summarizeText } from './execution-log'; import { createGatewayEgressGrant, restoreGatewaySandboxResult, revokeGatewayEgressGrant } from './egress-gateway-client'; import { refreshEgressGrantClaims } from './sandbox-egress'; import { buildSandboxExecuteRequest } from './sandbox-dispatch'; +import { getSandboxBackend, SandboxBackendError } 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'); } @@ -81,32 +76,29 @@ 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 responseRaw = await getSandboxBackend().execute( + { body: sandboxRequest.body, headers: sandboxRequest.headers }, + { + executionId: job.data.executionId ?? '', + language, + isSynthetic: isSyntheticJob, + signal: controller.signal, + tenantId: job.data.tenantId, + canonicalUserId: job.data.canonicalUserId, + runtimeSessionId: job.data.runtimeSessionId, + runtimeSessionMode: env.RUNTIME_SESSION_MODE, + }, + ); const responseData = egressGrantTokenForRestore ? await restoreGatewaySandboxResult({ grantId: egressGrantId, egressGrantToken: egressGrantTokenForRestore, - result: response.data, + result: responseRaw, isSynthetic: isSyntheticJob, signal: controller.signal, }) - : response.data; + : responseRaw; if (!isSyntheticJob) { logger.info('Sandbox response', summarizeSandboxResponse(responseData)); @@ -154,7 +146,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { const errorDetails = getAxiosErrorDetails(error); logger.error('Error processing job', errorDetails); - if (isAbortError(error)) { + if (error instanceof SandboxBackendError) { + 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 */