feat: optional AWS Lambda MicroVM stateful sandbox backend

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)
 # ============================================================================

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');
+  });
+});

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/<hook>` 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;

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;

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(

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}"

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) => {