From 7619a399edca09c4d31777e4c9fb52716b965c43 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 08:22:02 -0400 Subject: [PATCH 01/24] refactor: extract SandboxBackend seam from worker sandbox dispatch Moves the sandbox execute POST from workers.ts into a pluggable HttpSandboxBackend behind getSandboxBackend(). Adds CODEAPI_SANDBOX_BACKEND config (http default) and a startup policy check that rejects lambda-microvm until that backend lands. No wire behavior change: the signed request body and headers pass through byte-identical, and axios errors are rethrown untouched so the worker's abort/timeout/sandbox-error mapping is unchanged. --- service/src/config.ts | 10 ++ service/src/lifecycle.ts | 5 +- service/src/sandbox-backend/http.test.ts | 149 ++++++++++++++++++++++ service/src/sandbox-backend/http.ts | 34 +++++ service/src/sandbox-backend/index.test.ts | 37 ++++++ service/src/sandbox-backend/index.ts | 24 ++++ service/src/sandbox-backend/types.ts | 37 ++++++ service/src/secure-startup.test.ts | 18 +++ service/src/secure-startup.ts | 11 ++ service/src/workers.ts | 43 +++---- 10 files changed, 341 insertions(+), 27 deletions(-) create mode 100644 service/src/sandbox-backend/http.test.ts create mode 100644 service/src/sandbox-backend/http.ts create mode 100644 service/src/sandbox-backend/index.test.ts create mode 100644 service/src/sandbox-backend/index.ts create mode 100644 service/src/sandbox-backend/types.ts diff --git a/service/src/config.ts b/service/src/config.ts index 5abc57c..16afead 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -142,6 +142,16 @@ 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', }; const default_run_memory_limit = 256 * 1024 * 1024; diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 6deec64..9b21d05 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,7 @@ function setupQueueListeners(queue: Queue, name: string): void { export async function startupApiOnly(): Promise { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); + validateSandboxBackendPolicy(); await validateLifecycleAuthConfig(); // Set up queue listeners for monitoring (optional, for observability) @@ -92,6 +93,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 +127,7 @@ async function gracefulStartup(): Promise { logger.info('Starting up service (combined API + Workers)...'); validateApiHardenedConfig(); validateWorkerHardenedConfig(); + validateSandboxBackendPolicy(); await validateLifecycleAuthConfig(); try { 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..ba8c6d0 --- /dev/null +++ b/service/src/sandbox-backend/index.test.ts @@ -0,0 +1,37 @@ +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('rejects lambda-microvm until the backend lands', () => { + env.SANDBOX_BACKEND = 'lambda-microvm'; + expect(() => getSandboxBackend()).toThrow('lambda-microvm is not yet available'); + }); + + 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..e1c4e4d --- /dev/null +++ b/service/src/sandbox-backend/index.ts @@ -0,0 +1,24 @@ +import type { SandboxBackend } from './types'; +import { HttpSandboxBackend } from './http'; +import { env } from '../config'; + +export type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +export { HttpSandboxBackend } from './http'; + +let backend: SandboxBackend | undefined; + +function createBackend(): SandboxBackend { + if (env.SANDBOX_BACKEND === 'lambda-microvm') { + throw new Error('CODEAPI_SANDBOX_BACKEND=lambda-microvm is not yet available'); + } + 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/types.ts b/service/src/sandbox-backend/types.ts new file mode 100644 index 0000000..a33fc97 --- /dev/null +++ b/service/src/sandbox-backend/types.ts @@ -0,0 +1,37 @@ +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; +} diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 15224ad..a663c82 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -3,12 +3,14 @@ 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, gatewayUrl: env.EGRESS_GATEWAY_URL, grantSecret: env.EGRESS_GRANT_SECRET, privateKey: env.EXECUTION_MANIFEST_PRIVATE_KEY, @@ -24,6 +26,7 @@ function restore(): void { } Object.assign(process.env, savedEnv); env.HARDENED_SANDBOX_MODE = saved.hardened; + env.SANDBOX_BACKEND = saved.sandboxBackend; env.EGRESS_GATEWAY_URL = saved.gatewayUrl; env.EGRESS_GRANT_SECRET = saved.grantSecret; env.EXECUTION_MANIFEST_PRIVATE_KEY = saved.privateKey; @@ -120,3 +123,18 @@ describe('hardened CodeAPI startup config', () => { expect(() => validateEgressGatewayHardenedConfig()).toThrow('CODEAPI_EGRESS_LEDGER_REQUIRED'); }); }); + +describe('sandbox backend policy', () => { + test('accepts the default http backend', () => { + env.SANDBOX_BACKEND = 'http'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('rejects lambda-microvm regardless of hardened mode', () => { + env.SANDBOX_BACKEND = 'lambda-microvm'; + env.HARDENED_SANDBOX_MODE = false; + expect(() => validateSandboxBackendPolicy()).toThrow('lambda-microvm is not yet available'); + env.HARDENED_SANDBOX_MODE = true; + expect(() => validateSandboxBackendPolicy()).toThrow('lambda-microvm is not yet available'); + }); +}); diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 1d84e90..265323d 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -48,6 +48,17 @@ 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.SANDBOX_BACKEND !== 'lambda-microvm') return; + throw new SecureStartupConfigError( + 'CODEAPI_SANDBOX_BACKEND=lambda-microvm is not yet available; unset it or use "http"', + ); +} + 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/workers.ts b/service/src/workers.ts index 177e8fd..b860ca8 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 } 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,28 @@ 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, + runtimeSessionMode: 'stateless', + }, + ); 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)); From 63aacf64559276925debab33aeb8835a04f0958c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 08:29:16 -0400 Subject: [PATCH 02/24] feat: derive runtime session ids and add Redis session registry (dark launch) Adds runtime_session_hint to /exec (validated, optional) and derives rt_ server-side so a client hint can never collide across tenants. The id rides JobData into the sandbox backend context; stateless mode (default) derives nothing and enqueues byte-identical job data. The registry maps runtime_session_id -> MicroVM record in Redis with the replay-state lock discipline: SET NX PX mutex, CAS-delete release, token-fenced record writes/removals, monotonic generation counter for launch fencing, and a last-seen zset for the idle sweeper. No consumers yet - the Lambda backend lands behind CODEAPI_SANDBOX_BACKEND. --- service/src/config.ts | 13 ++ service/src/runtime-session/id.test.ts | 68 ++++++ service/src/runtime-session/id.ts | 59 ++++++ service/src/runtime-session/registry.test.ts | 150 +++++++++++++ service/src/runtime-session/registry.ts | 211 +++++++++++++++++++ service/src/service/router.ts | 17 ++ service/src/types/service.ts | 9 + service/src/workers.ts | 3 +- 8 files changed, 529 insertions(+), 1 deletion(-) create mode 100644 service/src/runtime-session/id.test.ts create mode 100644 service/src/runtime-session/id.ts create mode 100644 service/src/runtime-session/registry.test.ts create mode 100644 service/src/runtime-session/registry.ts diff --git a/service/src/config.ts b/service/src/config.ts index 16afead..34c8d1d 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -152,6 +152,19 @@ export const env = { 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, }; const default_run_memory_limit = 256 * 1024 * 1024; 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/registry.test.ts b/service/src/runtime-session/registry.test.ts new file mode 100644 index 0000000..8955813 --- /dev/null +++ b/service/src/runtime-session/registry.test.ts @@ -0,0 +1,150 @@ +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('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..7a40a2b --- /dev/null +++ b/service/src/runtime-session/registry.ts @@ -0,0 +1,211 @@ +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; + state: RuntimeSessionState; + generation: number; + launched_at?: number; + last_seen_at: number; + hard_deadline_at?: number; + workspace_checkpoint?: string; + last_error?: string; +} + +const SESS_PREFIX = 'rtsx:sess:'; +const LOCK_PREFIX = 'rtsx:lock:'; +const GEN_PREFIX = 'rtsx:gen:'; +const ACTIVE_ZSET = 'rtsx:active'; + +/** Launch budget placeholder until the Lambda backend config lands; the lock + * must outlive lock-wait + launch + execute so a live holder is never fenced + * mid-operation. */ +const LAUNCH_BUDGET_MS = 60_000; +export const RUNTIME_SESSION_LOCK_TTL_MS = env.JOB_TIMEOUT + LAUNCH_BUDGET_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; + 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('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 }); + } +} + +export async function readRuntimeSessionRecord(runtimeSessionId: string): Promise { + const data = await redis.get(`${SESS_PREFIX}${runtimeSessionId}`); + return data != null ? (JSON.parse(data) as RuntimeSessionRecord) : 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; +} + +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/service/router.ts b/service/src/service/router.ts index 9d9fd37..aa57646 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,21 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) return res.status(400).json({ error: `Unsupported language: ${rawLang}` }); } + let runtimeSessionId: string | undefined; + try { + runtimeSessionId = resolveRuntimeSessionIdForRequest({ + mode: env.RUNTIME_SESSION_MODE, + storageNamespace: identity.storageNamespace, + canonicalUserId: identity.canonicalUserId, + hint: validateRuntimeSessionHint(body.runtime_session_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 +235,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/workers.ts b/service/src/workers.ts index b860ca8..f05e3f6 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -85,7 +85,8 @@ async function processJobInner(job: t.ExecuteJob): Promise { signal: controller.signal, tenantId: job.data.tenantId, canonicalUserId: job.data.canonicalUserId, - runtimeSessionMode: 'stateless', + runtimeSessionId: job.data.runtimeSessionId, + runtimeSessionMode: env.RUNTIME_SESSION_MODE, }, ); From 9f7ee1a2d30faa5efaf2beffd95d6d63d92497b6 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 08:35:29 -0400 Subject: [PATCH 03/24] feat: Lambda MicroVM client wrapper, test fake, and op throttle Adds the LambdaMicrovmClient interface plus AwsLambdaMicrovmClient (@aws-sdk/client-lambda-microvms, isolated to lambda-client-aws.ts and absent from http-only bundles), a transport-free in-memory fake for bun tests, and Redis-backed per-second token buckets with poison backoff for the account-wide control-plane TPS limits. Mapping notes from the SDK typings: RunMicrovm takes imageIdentifier + imageVersion, connector ARN arrays, native idlePolicy (auto-suspend / auto-terminate / auto-resume), runHookPayload, and a clientToken idempotency key; auth tokens come back as an X-aws-proxy-auth header map with minute-granularity expiry (max 60). --- service/bun.lock | 47 +++++ service/package.json | 1 + .../runtime-session/lambda-client-aws.test.ts | 173 ++++++++++++++++ .../src/runtime-session/lambda-client-aws.ts | 163 +++++++++++++++ .../lambda-client-fake.test.ts | 76 +++++++ .../src/runtime-session/lambda-client-fake.ts | 192 ++++++++++++++++++ service/src/runtime-session/lambda-client.ts | 86 ++++++++ service/src/runtime-session/throttle.test.ts | 89 ++++++++ service/src/runtime-session/throttle.ts | 87 ++++++++ 9 files changed, 914 insertions(+) create mode 100644 service/src/runtime-session/lambda-client-aws.test.ts create mode 100644 service/src/runtime-session/lambda-client-aws.ts create mode 100644 service/src/runtime-session/lambda-client-fake.test.ts create mode 100644 service/src/runtime-session/lambda-client-fake.ts create mode 100644 service/src/runtime-session/lambda-client.ts create mode 100644 service/src/runtime-session/throttle.test.ts create mode 100644 service/src/runtime-session/throttle.ts 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/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..ffb0a17 --- /dev/null +++ b/service/src/runtime-session/lambda-client-aws.ts @@ -0,0 +1,163 @@ +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 { + 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, + 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..e5a966d --- /dev/null +++ b/service/src/runtime-session/lambda-client.ts @@ -0,0 +1,86 @@ +/** + * 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; + /** 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'; + +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/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); +} From 4d8d7e09c5489d99dceea60e64d98f4120a107f9 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 09:56:09 -0400 Subject: [PATCH 04/24] feat: Lambda MicroVM image target, lifecycle hooks, and stateless backend Adds the AWS Lambda MicroVM execution path (opt-in, config-gated behind CODEAPI_SANDBOX_BACKEND=lambda-microvm; default http is unchanged): - api/Dockerfile lambda-microvm-runner target: the existing sandbox-runner packaged for Lambda MicroVMs (arm64, port 8080, /pkgs baked, no libkrun launcher since the MicroVM is the VM boundary). - api lifecycle hooks at /aws/lambda-microvms/runtime/v1/{ready,run,resume, suspend,terminate}: no-op acks in Phase 1-2, /run captures the per-VM runHookPayload (Phase 3 checkpoint attachment point). - LambdaMicrovmSandboxBackend (stateless mode): run -> poll RUNNING -> mint X-aws-proxy-auth token -> health -> proxy /api/v2/execute -> terminate (incl. terminate-on-abort); throttle-aware, metrics-wired. - Startup policy rules (blocking-PTC reject, image-ARN required, hardened egress-connector required, token-TTL cap, no shell in prod). - entrypoint raises RLIMIT_NOFILE hard cap to 65536: the AL2023 MicroVM base caps it at 1024 (below the 2048 sandbox default), which made every in-guest nsjail job fail setrlimit with EPERM. Verified on live AWS. Live-verified on AWS Lambda MicroVMs (us-east-1): image builds and snapshots, nsjail runs with additionalOsCapabilities ALL, bash+python execute (code 0), and suspend/resume preserves process+memory state with ~1.2s auto-resume-on-traffic. --- api/Dockerfile | 20 ++ api/src/api/lifecycle.test.ts | 46 ++++ api/src/api/lifecycle.ts | 83 ++++++ api/src/entrypoint.sh | 9 + api/src/index.ts | 2 + scripts/build-lambda-microvm-artifact.sh | 104 ++++++++ service/src/config.ts | 27 ++ service/src/metrics.ts | 24 ++ service/src/sandbox-backend/index.test.ts | 7 +- service/src/sandbox-backend/index.ts | 25 +- .../sandbox-backend/lambda-microvm.test.ts | 245 ++++++++++++++++++ service/src/sandbox-backend/lambda-microvm.ts | 227 ++++++++++++++++ service/src/sandbox-backend/types.ts | 22 ++ service/src/secure-startup.test.ts | 76 +++++- service/src/secure-startup.ts | 35 ++- service/src/workers.ts | 6 +- 16 files changed, 945 insertions(+), 13 deletions(-) create mode 100644 api/src/api/lifecycle.test.ts create mode 100644 api/src/api/lifecycle.ts create mode 100755 scripts/build-lambda-microvm-artifact.sh create mode 100644 service/src/sandbox-backend/lambda-microvm.test.ts create mode 100644 service/src/sandbox-backend/lambda-microvm.ts diff --git a/api/Dockerfile b/api/Dockerfile index ff4e848..164e41b 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -148,6 +148,26 @@ 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 + +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..13a1861 --- /dev/null +++ b/api/src/api/lifecycle.ts @@ -0,0 +1,83 @@ +import express, { Router, type Request, type Response } from 'express'; +import { logger } from '../logger'; + +/** + * 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() }; + } 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', ackHook('terminate')); + +export default lifecycleRouter; diff --git a/api/src/entrypoint.sh b/api/src/entrypoint.sh index 69145b1..c7f576d 100755 --- a/api/src/entrypoint.sh +++ b/api/src/entrypoint.sh @@ -3,6 +3,15 @@ set -e echo "Starting NsJail sandbox API..." +# Raise the RLIMIT_NOFILE hard cap 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. No-op where the limit is already higher +# (e.g. Docker's ~1M default) or where CAP_SYS_RESOURCE is unavailable. +ulimit -Hn 65536 2>/dev/null || true +ulimit -Sn 65536 2>/dev/null || true + 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/scripts/build-lambda-microvm-artifact.sh b/scripts/build-lambda-microvm-artifact.sh new file mode 100755 index 0000000..c81816e --- /dev/null +++ b/scripts/build-lambda-microvm-artifact.sh @@ -0,0 +1,104 @@ +#!/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 < \\ + --hook-port 8080 \\ + --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. +- /ready + /run/resume/suspend/terminate hooks are served on port 8080 at + /aws/lambda-microvms/runtime/v1/*. +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/src/config.ts b/service/src/config.ts index 34c8d1d..5882726 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() === '') { @@ -165,6 +171,27 @@ export const env = { ? 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, + 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, + LAMBDA_MICROVM_ALLOW_SHELL: process.env.LAMBDA_MICROVM_ALLOW_SHELL === 'true', }; const default_run_memory_limit = 256 * 1024 * 1024; diff --git a/service/src/metrics.ts b/service/src/metrics.ts index a6ce6cd..522456a 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -194,6 +194,30 @@ 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, +}); + // -- 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/sandbox-backend/index.test.ts b/service/src/sandbox-backend/index.test.ts index ba8c6d0..3061348 100644 --- a/service/src/sandbox-backend/index.test.ts +++ b/service/src/sandbox-backend/index.test.ts @@ -19,9 +19,12 @@ describe('getSandboxBackend', () => { expect(getSandboxBackend()).toBe(backend); }); - test('rejects lambda-microvm until the backend lands', () => { + test('selects the lambda-microvm backend when configured', async () => { env.SANDBOX_BACKEND = 'lambda-microvm'; - expect(() => getSandboxBackend()).toThrow('lambda-microvm is not yet available'); + 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', () => { diff --git a/service/src/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts index e1c4e4d..60afec4 100644 --- a/service/src/sandbox-backend/index.ts +++ b/service/src/sandbox-backend/index.ts @@ -1,15 +1,38 @@ import type { SandboxBackend } from './types'; +import { LambdaMicrovmSandboxBackend } from './lambda-microvm'; import { HttpSandboxBackend } from './http'; 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') { - throw new Error('CODEAPI_SANDBOX_BACKEND=lambda-microvm is not yet available'); + 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, + 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, + jobTimeoutMs: env.JOB_TIMEOUT, + }, + }); } return new HttpSandboxBackend(); } 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..00086a3 --- /dev/null +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -0,0 +1,245 @@ +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 { 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 mock: InstanceType; + +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: 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); + captured = []; + healthStatus = 200; + executeDelayMs = 0; +}); + +afterEach(() => { + resetThrottleRedis(); +}); + +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, + jobTimeoutMs: 300_000, + ...overrides, + }; +} + +function makeBackend(fake: FakeLambdaMicrovmClient, cfg?: Partial): LambdaMicrovmSandboxBackend { + return new LambdaMicrovmSandboxBackend({ + clientFactory: () => Promise.resolve(fake), + config: config(cfg), + pollIntervalMs: 5, + }); +} + +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); + }); +}); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts new file mode 100644 index 0000000..7c2ebeb --- /dev/null +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -0,0 +1,227 @@ +import axios from 'axios'; +import { nanoid } from 'nanoid'; +import type { LambdaMicrovmClient, MicrovmDescription } from '../runtime-session/lambda-client'; +import type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +import { LambdaMicrovmApiError } from '../runtime-session/lambda-client'; +import { MicrovmOpThrottledError, acquireOpBudget, poisonOpBucket } from '../runtime-session/throttle'; +import { microvmLaunches, microvmLaunchDuration, microvmTerminations, microvmThrottleEvents } from '../metrics'; +import { injectTraceHeaders, withSpan } from '../telemetry'; +import { SandboxBackendError } from './types'; +import { Jobs } from '../enum'; +import logger from '../logger'; + +export interface LambdaMicrovmBackendConfig { + imageArn: string; + imageVersion?: string; + executionRoleArn?: string; + ingressConnectorArns?: string[]; + egressConnectorArns?: string[]; + port: number; + maxDurationSeconds: number; + authTokenTtlSeconds: number; + launchTimeoutMs: number; + healthTimeoutMs: number; + launchTps: number; + jobTimeoutMs: number; +} + +interface LambdaMicrovmBackendDeps { + clientFactory: () => Promise; + config: LambdaMicrovmBackendConfig; + pollIntervalMs?: number; +} + +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(/\/+$/, '')}`; +} + +/** + * Stateless Lambda MicroVM backend: one VM per execution + * (run -> poll RUNNING -> health -> execute -> terminate). Runtime-session + * reuse (find-or-launch on the registry) lands in the next phase; the + * startup policy rejects non-stateless modes until then. + */ +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; + + constructor(deps: LambdaMicrovmBackendDeps) { + this.clientFactory = deps.clientFactory; + this.config = deps.config; + this.pollIntervalMs = deps.pollIntervalMs ?? 500; + } + + private client(): Promise { + this.clientPromise ??= this.clientFactory(); + return this.clientPromise; + } + + async execute(req: SandboxTransportRequest, ctx: SandboxExecuteContext): Promise { + const client = await this.client(); + const vm = await this.launch(client, ctx); + let terminateReason = 'stateless'; + try { + const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + const token = await client.createMicrovmAuthToken({ + microvmId: vm.microvmId, + port: this.config.port, + ttlSeconds: this.config.authTokenTtlSeconds, + }); + await this.assertHealthy(base, token.token, ctx); + + return await 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, + }, + signal: ctx.signal, + }, + ); + if (response.status !== 200) { + throw new Error('Error from sandbox'); + } + return response.data; + }, 'CLIENT'); + } catch (error) { + terminateReason = ctx.signal.aborted ? 'timeout' : 'error'; + throw error; + } finally { + await this.terminate(client, vm.microvmId, terminateReason); + } + } + + private async launch(client: LambdaMicrovmClient, ctx: SandboxExecuteContext): 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 { + /* Stateless 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, + ); + vm = await client.runMicrovm({ + imageIdentifier: this.config.imageArn, + imageVersion: this.config.imageVersion, + executionRoleArn: this.config.executionRoleArn, + ingressConnectorArns: this.config.ingressConnectorArns, + egressConnectorArns: this.config.egressConnectorArns, + maximumDurationSeconds: maxDurationSeconds, + clientToken: ctx.executionId !== '' ? `exec-${ctx.executionId}` : `exec-${nanoid()}`, + }); + } 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'); + throw 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 }, + 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 index a33fc97..190ab70 100644 --- a/service/src/sandbox-backend/types.ts +++ b/service/src/sandbox-backend/types.ts @@ -35,3 +35,25 @@ export interface SandboxBackend { 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 a663c82..a3baf2b 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -11,6 +11,12 @@ 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, @@ -27,6 +33,12 @@ 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; @@ -125,16 +137,70 @@ describe('hardened CodeAPI startup config', () => { }); 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; + } + test('accepts the default http backend', () => { env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'stateless'; expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); - test('rejects lambda-microvm regardless of hardened mode', () => { - env.SANDBOX_BACKEND = 'lambda-microvm'; - env.HARDENED_SANDBOX_MODE = false; - expect(() => validateSandboxBackendPolicy()).toThrow('lambda-microvm is not yet available'); + 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('rejects non-stateless session modes until orchestration lands', () => { + configureValidLambda(); + env.RUNTIME_SESSION_MODE = 'affinity'; + expect(() => validateSandboxBackendPolicy()).toThrow('session orchestration is not yet available'); + }); + + 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; - expect(() => validateSandboxBackendPolicy()).toThrow('lambda-microvm is not yet available'); + 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 265323d..337ad81 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -53,10 +53,39 @@ export function validateWorkerHardenedConfig(): void { * 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', + ); + } if (env.SANDBOX_BACKEND !== 'lambda-microvm') return; - throw new SecureStartupConfigError( - 'CODEAPI_SANDBOX_BACKEND=lambda-microvm is not yet available; unset it or use "http"', - ); + + if (env.PTC_MODE === 'blocking') { + throw new SecureStartupConfigError( + 'PTC replay is the only supported PTC mode for the lambda-microvm backend (unset PTC_MODE=blocking)', + ); + } + if (env.LAMBDA_MICROVM_IMAGE_ARN.trim().length === 0) { + throw new SecureStartupConfigError('LAMBDA_MICROVM_IMAGE_ARN is required for the lambda-microvm backend'); + } + if (env.RUNTIME_SESSION_MODE !== 'stateless') { + throw new SecureStartupConfigError( + 'Runtime session orchestration is not yet available for the lambda-microvm backend; set CODEAPI_RUNTIME_SESSION_MODE=stateless', + ); + } + 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', + ); + } } export function validateEgressGatewayHardenedConfig(): void { diff --git a/service/src/workers.ts b/service/src/workers.ts index f05e3f6..f15c393 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -10,7 +10,7 @@ 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 } from './sandbox-backend'; +import { getSandboxBackend, SandboxBackendError } from './sandbox-backend'; import { isSyntheticPrincipalSource } from './auth/synthetic'; import { withSpan, withTraceContext } from './telemetry'; import logger from './logger'; @@ -146,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 */ From ac4fdef0f72bd82261e753a98fbe31962957fa6c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 11:45:04 -0400 Subject: [PATCH 05/24] feat: persistent session workspace for stateful MicroVM sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Turns the semi-stateless runner stateful when a VM is bound to a runtime session: instead of a fresh /mnt/data workspace per /execute, the VM reuses ONE persistent workspace and one pinned UID across calls, so files, installed packages, and chDB dirs survive between tool calls. Modular and off by default: gated by TWO independent locks, so the legacy fresh-per-job path is byte-identical when either is absent. 1. Image-level SANDBOX_SESSION_WORKSPACE_ENABLED — set only in the lambda-microvm-runner target; the K8s image cannot enter session mode regardless of any payload. 2. Per-launch /run runHookPayload {runtime_session_id, session_workspace} — the control plane opts a specific VM in. Mechanism: - session-workspace.ts: process singleton bound by the /run hook, unbound by /terminate; holds the pinned UID + output-diff and priming-dedup state. - workspace-isolation.ts: ensureSessionWorkspace (stable id, contents preserved, reaper-protected) + resetSessionWorkspace teardown. - Job branches only at three seams — prime (reuse workspace+UID, skip re-downloading unchanged inputs), walkDir file surfacing (skip prior outputs unchanged by size+mtime — output diffing), cleanup (keep workspace+UID for the session). Verified end-to-end in the arm64 runner container: fresh mode wipes between calls; session mode persists files across calls (incl. across languages) and accumulates; /terminate wipes and rebind gives a clean slate. 274 api tests pass (fresh path unchanged). --- api/Dockerfile | 3 +- api/src/api/lifecycle.ts | 13 ++- api/src/api/v2.ts | 2 + api/src/config.ts | 6 ++ api/src/job.ts | 80 ++++++++++++++- api/src/session-workspace.test.ts | 97 ++++++++++++++++++ api/src/session-workspace.ts | 147 ++++++++++++++++++++++++++++ api/src/workspace-isolation.test.ts | 45 +++++++++ api/src/workspace-isolation.ts | 51 ++++++++++ 9 files changed, 438 insertions(+), 6 deletions(-) create mode 100644 api/src/session-workspace.test.ts create mode 100644 api/src/session-workspace.ts diff --git a/api/Dockerfile b/api/Dockerfile index 164e41b..b3ee5b0 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -163,7 +163,8 @@ FROM sandbox-build AS lambda-microvm-runner COPY --from=package-builder /pkgs /pkgs ENV PORT=8080 \ - SANDBOX_PACKAGES_DIRECTORY=/pkgs + SANDBOX_PACKAGES_DIRECTORY=/pkgs \ + SANDBOX_SESSION_WORKSPACE_ENABLED=true EXPOSE 8080/tcp ENTRYPOINT ["/sandbox_api/entrypoint.sh"] diff --git a/api/src/api/lifecycle.ts b/api/src/api/lifecycle.ts index 13a1861..2267efb 100644 --- a/api/src/api/lifecycle.ts +++ b/api/src/api/lifecycle.ts @@ -1,5 +1,6 @@ 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 @@ -46,6 +47,11 @@ export function applyRunHook(body: unknown): MicrovmRunContext { 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 }, @@ -78,6 +84,11 @@ lifecycleRouter.post('/run', express.json({ limit: '32kb' }), (req: Request, res lifecycleRouter.post('/resume', ackHook('resume')); lifecycleRouter.post('/suspend', ackHook('suspend')); -lifecycleRouter.post('/terminate', ackHook('terminate')); + +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..af643c3 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -12,6 +12,7 @@ import { activeSandboxExecutions, recordSandboxExecution } from '../metrics'; import { classifySandboxSafeError } from '../safe-error'; import { withSpan } from '../telemetry'; import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; +import { getBoundSessionWorkspace } from '../session-workspace'; const router = express.Router(); const SYNTHETIC_PRINCIPAL_SOURCE = 'synthetic_test'; @@ -194,6 +195,7 @@ function getJob( egress_grant: egressGrantToken, tool_call_socket_enabled: toolCallSocketEnabled, is_synthetic: isSynthetic, + session: getBoundSessionWorkspace() ?? null, }); } 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/job.ts b/api/src/job.ts index 291d76d..3839b44 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, @@ -680,6 +681,10 @@ export class Job { 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 +703,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 }); @@ -792,8 +800,13 @@ export class Job { } 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 +816,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 +829,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 +837,38 @@ 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) this.session.markPrimed(name, file.id); + } + + 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; + const hash = 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; } @@ -1478,11 +1524,14 @@ export class Job { } let size: number; + let mtimeMs: number; try { /* 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; + mtimeMs = st.mtimeMs; } catch (err) { this.log.debug({ path: relativePath, err }, 'walkDir: unable to stat file'); return { collected: false, truncated: false, stopLoop: false }; @@ -1491,6 +1540,17 @@ export class Job { return { collected: false, truncated: false, stopLoop: false }; } + /* Session mode output diffing: a persistent workspace still holds prior + * jobs' outputs. Skip any file already surfaced to the client whose + * size+mtime are unchanged, so only genuine deltas are re-uploaded. An + * input file (tracked in inputFileHashes) is left to the existing + * modified/unchanged classification below. */ + const outputSignature = `${size}:${mtimeMs}`; + if (this.session && !this.inputFileHashes.has(relativePath) + && this.session.isSurfaced(relativePath, outputSignature)) { + return { collected: false, truncated: false, stopLoop: false }; + } + const inputFileInfo = this.inputFileHashes.get(relativePath); const existingFile = inputByName.get(relativePath); let wasModified = false; @@ -1528,6 +1588,7 @@ export class Job { } this.sessionFiles.push(fileData); this.generatedFiles.push({ id: newId, name: relativePath, path: fullPath }); + if (this.session) this.session.markSurfaced(relativePath, outputSignature); return { collected: true, truncated: false, stopLoop: false }; } @@ -1818,6 +1879,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-workspace.test.ts b/api/src/session-workspace.test.ts new file mode 100644 index 0000000..86d09d1 --- /dev/null +++ b/api/src/session-workspace.test.ts @@ -0,0 +1,97 @@ +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, + 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(); + }); + + 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(); + ws.markPrimed('in.csv', 'file_abc'); + expect(ws.primedInputId('in.csv')).toBe('file_abc'); + + 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 }); + } + }); +}); diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts new file mode 100644 index 0000000..9711178 --- /dev/null +++ b/api/src/session-workspace.ts @@ -0,0 +1,147 @@ +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-launch `/run` runHookPayload opting in. When neither is + * active, `getBoundSessionWorkspace()` returns undefined and the runner falls + * back to the untouched fresh-per-job path. + */ + +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 }; +} + +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 -> storage file id already primed onto disk, so an unchanged + * input delivered again is not re-downloaded (priming dedup). */ + 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; + } + + 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); + } + + primedInputId(relPath: string): string | undefined { + return this.primed.get(relPath); + } + + markPrimed(relPath: string, storageFileId: string): void { + this.primed.set(relPath, storageFileId); + } + + /** 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); From f80700e314a25644c52f0670567c253f9840ea81 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 11:52:16 -0400 Subject: [PATCH 06/24] feat: backend session orchestration (find-or-launch + warm-VM reuse) Connects the proven runner session workspace to the product: when a runtime session id is present and the mode is affinity/strict, the Lambda backend finds-or-launches ONE warm VM per runtime_session_id via the registry, delivers the /run runHookPayload {runtime_session_id, session_workspace:true} that activates the runner's persistent workspace, and reuses that VM across executes. AWS idlePolicy (autoResume) suspends the VM when idle and wakes it on the next request, so there is no explicit resume in the hot path. - Serializes per session on the registry lock; strict contention -> 409 (publicExecutionFailure maps RUNTIME_SESSION_BUSY), affinity contention -> correct stateless one-shot fallback (files always ride the payload). - Generation-fenced launch: a fenced worker terminates its orphan VM. - Stateless path unchanged (one VM per exec, terminate after). - Lifts the stateless-only startup policy; adds session/fallback/lock metrics. 345 service tests pass, incl. reuse/serialize/fallback/fence. --- service/src/metrics.ts | 16 ++ service/src/sandbox-backend/index.ts | 3 + .../sandbox-backend/lambda-microvm.test.ts | 116 ++++++++ service/src/sandbox-backend/lambda-microvm.ts | 262 +++++++++++++++--- service/src/secure-startup.test.ts | 6 +- service/src/secure-startup.ts | 5 - service/src/utils.test.ts | 16 ++ service/src/utils.ts | 11 + 8 files changed, 382 insertions(+), 53 deletions(-) diff --git a/service/src/metrics.ts b/service/src/metrics.ts index 522456a..cf4c54e 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -218,6 +218,22 @@ export const microvmThrottleEvents = new Counter({ 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', +}); + // -- 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/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts index 60afec4..6161c1b 100644 --- a/service/src/sandbox-backend/index.ts +++ b/service/src/sandbox-backend/index.ts @@ -31,6 +31,9 @@ function createBackend(): SandboxBackend { healthTimeoutMs: env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS, launchTps: env.LAMBDA_MICROVM_LAUNCH_TPS, jobTimeoutMs: env.JOB_TIMEOUT, + idleSeconds: env.LAMBDA_MICROVM_IDLE_SECONDS, + suspendedSeconds: env.LAMBDA_MICROVM_SUSPEND_SECONDS, + lockWaitMs: env.RUNTIME_SESSION_LOCK_WAIT_MS, }, }); } diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 00086a3..1fd2939 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -7,6 +7,12 @@ import { resetRedisForTests as resetThrottleRedis, setRedisForTests as setThrottleRedis, } from '../runtime-session/throttle'; +import { + acquireRuntimeSessionLock, + readRuntimeSessionRecord, + resetRedisForTests as resetRegistryRedis, + setRedisForTests as setRegistryRedis, +} from '../runtime-session/registry'; import { LambdaMicrovmSandboxBackend, normalizeMicrovmEndpoint, type LambdaMicrovmBackendConfig } from './lambda-microvm'; import { SandboxBackendError } from './types'; import type { SandboxExecuteContext, SandboxTransportRequest } from './types'; @@ -70,6 +76,7 @@ beforeEach(async () => { mock = new RedisMock(); await mock.flushall(); setThrottleRedis(mock); + setRegistryRedis(mock); captured = []; healthStatus = 200; executeDelayMs = 0; @@ -77,6 +84,7 @@ beforeEach(async () => { afterEach(() => { resetThrottleRedis(); + resetRegistryRedis(); }); function config(overrides: Partial = {}): LambdaMicrovmBackendConfig { @@ -90,6 +98,9 @@ function config(overrides: Partial = {}): LambdaMicr healthTimeoutMs: 1_000, launchTps: 50, jobTimeoutMs: 300_000, + idleSeconds: 300, + suspendedSeconds: 1_800, + lockWaitMs: 500, ...overrides, }; } @@ -243,3 +254,108 @@ describe('LambdaMicrovmSandboxBackend stateless execution', () => { 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 session VM with the workspace runHookPayload + idlePolicy and records it', 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; + }; + expect(JSON.parse(runArgs.runHookPayload as string)).toEqual({ + runtime_session_id: 'rt_session_1', + session_workspace: true, + }); + expect(runArgs.idlePolicy?.autoResume).toBe(true); + expect(runArgs.clientToken).toBe('sess-rt_session_1-1'); + expect(runArgs.maximumDurationSeconds).toBe(28_800); + + 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('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(); + }); +}); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 7c2ebeb..857dbe7 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -1,15 +1,37 @@ import axios from 'axios'; import { nanoid } from 'nanoid'; -import type { LambdaMicrovmClient, MicrovmDescription } from '../runtime-session/lambda-client'; +import type { LambdaMicrovmClient, MicrovmDescription, MicrovmIdlePolicy } from '../runtime-session/lambda-client'; import type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; +import type { RuntimeSessionRecord } from '../runtime-session/registry'; import { LambdaMicrovmApiError } from '../runtime-session/lambda-client'; import { MicrovmOpThrottledError, acquireOpBudget, poisonOpBucket } from '../runtime-session/throttle'; -import { microvmLaunches, microvmLaunchDuration, microvmTerminations, microvmThrottleEvents } from '../metrics'; +import { + allocateRuntimeSessionGeneration, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + 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'; +/** Payload delivered to the MicroVM /run hook to activate the runner's + * persistent session workspace (see api/src/session-workspace.ts). */ +function sessionRunHookPayload(runtimeSessionId: string): string { + return JSON.stringify({ runtime_session_id: runtimeSessionId, session_workspace: true }); +} + export interface LambdaMicrovmBackendConfig { imageArn: string; imageVersion?: string; @@ -23,6 +45,10 @@ export interface LambdaMicrovmBackendConfig { healthTimeoutMs: number; launchTps: number; jobTimeoutMs: number; + /* Session-mode (find-or-launch) tuning. */ + idleSeconds: number; + suspendedSeconds: number; + lockWaitMs: number; } interface LambdaMicrovmBackendDeps { @@ -41,11 +67,23 @@ export function normalizeMicrovmEndpoint(endpoint: string): string { return `https://${endpoint.replace(/\/+$/, '')}`; } +interface LaunchOptions { + clientToken: string; + runHookPayload?: string; + idlePolicy?: MicrovmIdlePolicy; + maxDurationSeconds: number; +} + /** - * Stateless Lambda MicroVM backend: one VM per execution - * (run -> poll RUNNING -> health -> execute -> terminate). Runtime-session - * reuse (find-or-launch on the registry) lands in the next phase; the - * startup policy rejects non-stateless modes until then. + * 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, deliver the /run payload that + * activates the runner's persistent workspace, and reuse it 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; @@ -67,39 +105,30 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { async execute(req: SandboxTransportRequest, ctx: SandboxExecuteContext): Promise { const client = await this.client(); - const vm = await this.launch(client, ctx); + 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 { - const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); - const token = await client.createMicrovmAuthToken({ - microvmId: vm.microvmId, - port: this.config.port, - ttlSeconds: this.config.authTokenTtlSeconds, - }); - await this.assertHealthy(base, token.token, ctx); - - return await 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, - }, - signal: ctx.signal, - }, - ); - if (response.status !== 200) { - throw new Error('Error from sandbox'); - } - return response.data; - }, 'CLIENT'); + return await this.proxyExecute(client, vm, req, ctx); } catch (error) { terminateReason = ctx.signal.aborted ? 'timeout' : 'error'; throw error; @@ -108,7 +137,152 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { } } - private async launch(client: LambdaMicrovmClient, ctx: SandboxExecuteContext): Promise { + private async executeSession( + client: LambdaMicrovmClient, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + runtimeSessionId: string, + ): Promise { + 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); + } + + try { + const existing = await readRuntimeSessionRecord(runtimeSessionId); + const vm = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); + const result = await this.proxyExecute(client, vm, req, ctx); + /* 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(); + const settled = await readRuntimeSessionRecord(runtimeSessionId); + if (settled) { + await writeRuntimeSessionRecord({ ...settled, state: 'RUNNING', last_seen_at: now }, lockToken); + } + await touchRuntimeSessionActive(runtimeSessionId, now); + return result; + } finally { + await releaseRuntimeSessionLock(runtimeSessionId, lockToken); + } + } + + private async findOrLaunchSession( + client: LambdaMicrovmClient, + ctx: SandboxExecuteContext, + runtimeSessionId: string, + record: RuntimeSessionRecord | null, + lockToken: string, + ): Promise { + const deadlineHeadroomMs = this.config.jobTimeoutMs + 30_000; + const reusable = record + && record.state === 'RUNNING' + && record.microvm_id + && record.endpoint + && (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 { microvmId: record.microvm_id as string, state: 'RUNNING', endpoint: record.endpoint }; + } + + 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}`, + runHookPayload: sessionRunHookPayload(runtimeSessionId), + 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, + 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`); + } + return vm; + } + + private async proxyExecute( + client: LambdaMicrovmClient, + vm: MicrovmDescription, + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + ): Promise { + const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); + const token = await client.createMicrovmAuthToken({ + microvmId: vm.microvmId, + port: this.config.port, + ttlSeconds: this.config.authTokenTtlSeconds, + }); + await this.assertHealthy(base, token.token, ctx); + + 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, + }, + 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', { @@ -126,20 +300,16 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { let vm: MicrovmDescription; try { - /* Stateless 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, - ); vm = await client.runMicrovm({ imageIdentifier: this.config.imageArn, imageVersion: this.config.imageVersion, executionRoleArn: this.config.executionRoleArn, ingressConnectorArns: this.config.ingressConnectorArns, egressConnectorArns: this.config.egressConnectorArns, - maximumDurationSeconds: maxDurationSeconds, - clientToken: ctx.executionId !== '' ? `exec-${ctx.executionId}` : `exec-${nanoid()}`, + maximumDurationSeconds: opts.maxDurationSeconds, + runHookPayload: opts.runHookPayload, + idlePolicy: opts.idlePolicy, + clientToken: opts.clientToken, }); } catch (error) { if (error instanceof LambdaMicrovmApiError && error.kind === 'throttled') { diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index a3baf2b..39d3bcf 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -177,10 +177,12 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).toThrow('LAMBDA_MICROVM_IMAGE_ARN is required'); }); - test('rejects non-stateless session modes until orchestration lands', () => { + test('accepts affinity and strict session modes on the lambda backend', () => { configureValidLambda(); env.RUNTIME_SESSION_MODE = 'affinity'; - expect(() => validateSandboxBackendPolicy()).toThrow('session orchestration is not yet available'); + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + env.RUNTIME_SESSION_MODE = 'strict'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); test('hardened mode requires an egress connector', () => { diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 337ad81..218cb7e 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -68,11 +68,6 @@ export function validateSandboxBackendPolicy(): void { if (env.LAMBDA_MICROVM_IMAGE_ARN.trim().length === 0) { throw new SecureStartupConfigError('LAMBDA_MICROVM_IMAGE_ARN is required for the lambda-microvm backend'); } - if (env.RUNTIME_SESSION_MODE !== 'stateless') { - throw new SecureStartupConfigError( - 'Runtime session orchestration is not yet available for the lambda-microvm backend; set CODEAPI_RUNTIME_SESSION_MODE=stateless', - ); - } 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)', 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; From 60fdf3baa8af5b339b375dd212d1a67941b6da76 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 12:13:41 -0400 Subject: [PATCH 07/24] feat: session workspace checkpoint/restore for perceived statefulness across expiry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes an expiring/evicted MicroVM's state survive a relaunch — the difference between real statefulness and just re-implementing the existing file-ref system. The file-ref path only brings back files surfaced as CodeEnvRefs; checkpoint/restore brings back the WHOLE workspace: pip-installed packages, venvs, chDB dirs, caches, and files with unsupported extensions. Two runner endpoints (session-mode only, 409 otherwise so the legacy runner exposes nothing new): - GET /api/v2/session/checkpoint streams tar.gz of the workspace - POST /api/v2/session/restore replaces the workspace from one, re-owned to the session's pinned UID Control-plane driven: the orchestrator pulls the checkpoint over the authed proxy and owns the S3 write, so the untrusted VM never gets S3 credentials (matches the report's checkpoint-capability security model). Verified end-to-end with two containers simulating VM expiry: VM1 builds a python module tree + a 2KB unsupported-extension binary, is terminated; a fresh VM2 shows the state absent (all file-refs give you), then after restore imports the module (greet()=42) and reads the binary (2048 bytes) — full workspace continuity across a VM swap. 276 api tests pass. --- api/src/api/v2.ts | 11 ++++ api/src/session-checkpoint.test.ts | 38 ++++++++++++ api/src/session-checkpoint.ts | 95 ++++++++++++++++++++++++++++++ api/src/session-workspace.ts | 8 +++ 4 files changed, 152 insertions(+) create mode 100644 api/src/session-checkpoint.test.ts create mode 100644 api/src/session-checkpoint.ts diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index af643c3..4574153 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -13,6 +13,7 @@ import { classifySandboxSafeError } from '../safe-error'; import { withSpan } from '../telemetry'; import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; import { getBoundSessionWorkspace } from '../session-workspace'; +import { streamSessionCheckpoint, restoreSessionCheckpoint } from '../session-checkpoint'; const router = express.Router(); const SYNTHETIC_PRINCIPAL_SOURCE = 'synthetic_test'; @@ -240,6 +241,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' }); } @@ -443,4 +446,12 @@ 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. */ +router.get('/session/checkpoint', (_req: Request, res: Response) => streamSessionCheckpoint(res)); +router.post('/session/restore', (req: Request, res: Response) => restoreSessionCheckpoint(req, res)); + export default router; 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..c129337 --- /dev/null +++ b/api/src/session-checkpoint.ts @@ -0,0 +1,95 @@ +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 { 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; + } + await session.ownership(); + + 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(); + } +} + +/** 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), so extraction stays within the workspace root without needing + * GNU-only flags. Production hardening: verify/scan the archive before + * trusting a restore from shared storage. */ + const tar = spawn('tar', ['-xzf', '-', '-C', SANDBOX_WORKSPACE_ROOT], { + 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); + res.status(200).json({ status: 'restored', dir: path.basename(dir) }); + } catch (error) { + logger.error({ err: error }, 'Failed to restore session checkpoint'); + if (!res.headersSent) res.status(500).json({ message: 'restore failed' }); + } +} + +async function chownRecursive(dir: string, uid: number, gid: number): Promise { + await fsp.chown(dir, uid, gid).catch(() => {}); + const entries = await fsp.readdir(dir, { withFileTypes: true }); + for (const entry of entries) { + const full = path.join(dir, entry.name); + await fsp.chown(full, uid, gid).catch(() => {}); + if (entry.isDirectory()) await chownRecursive(full, uid, gid); + } +} diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index 9711178..66726ed 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -82,6 +82,14 @@ export class SessionWorkspace { 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; } From a38933ab0647bcd2bb58e16f4da237db817a7db3 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 6 Jul 2026 09:36:52 -0400 Subject: [PATCH 08/24] feat: auto-checkpoint session workspaces to S3 with restore on relaunch Closes the 8h-rollover / eviction story so perceived statefulness is automatic instead of control-plane-by-hand. After each successful session exec (lock still held), pull the workspace tar from the warm VM and store it to S3 under a deterministic key (rtsx-checkpoints/.tar.gz) so recovery survives even registry loss; record the pointer under the same fenced write. On relaunch, a fresh session VM restores its predecessor's checkpoint before the first exec, making an expired/evicted VM invisible. - Coverage is complete and tear-free: the workspace only mutates during an exec and execs serialize on the session lock, so the post-exec checkpoint always captures the latest state; a busy lock means a newer exec will checkpoint instead. - Never fatal: a missed checkpoint degrades to file-ref recovery, a failed restore to a fresh workspace ('relaunched must be correct'). - Off => warm reuse still works, cross-VM restore falls back to file refs. CheckpointStore injected (Minio prod / Memory in tests); byte cap + timeout bound the transfer; checkpoint/restore/bytes metrics. 354 service tests pass incl. checkpoint-after-exec, restore-before-first -exec ordering, no-restore-on-reuse, disabled skip, failure non-fatal. --- service/src/config.ts | 7 + service/src/metrics.ts | 21 +++ .../runtime-session/checkpoint-store.test.ts | 34 ++++ .../src/runtime-session/checkpoint-store.ts | 72 ++++++++ service/src/runtime-session/checkpoint.ts | 167 ++++++++++++++++++ service/src/runtime-session/registry.ts | 1 + service/src/sandbox-backend/index.ts | 11 ++ .../sandbox-backend/lambda-microvm.test.ts | 107 ++++++++++- service/src/sandbox-backend/lambda-microvm.ts | 69 +++++++- 9 files changed, 486 insertions(+), 3 deletions(-) create mode 100644 service/src/runtime-session/checkpoint-store.test.ts create mode 100644 service/src/runtime-session/checkpoint-store.ts create mode 100644 service/src/runtime-session/checkpoint.ts diff --git a/service/src/config.ts b/service/src/config.ts index 5882726..7814119 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -192,6 +192,13 @@ export const env = { LAMBDA_MICROVM_RESUME_TPS: Number(process.env.LAMBDA_MICROVM_RESUME_TPS) || 4, LAMBDA_MICROVM_SUSPEND_TPS: Number(process.env.LAMBDA_MICROVM_SUSPEND_TPS) || 1, 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/metrics.ts b/service/src/metrics.ts index cf4c54e..0175cca 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -234,6 +234,27 @@ export const microvmActiveSessions = new Gauge({ 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..cd261fc --- /dev/null +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, test } from 'bun:test'; +import { MemoryCheckpointStore, checkpointObjectKey } from './checkpoint-store'; + +describe('checkpoint store', () => { + test('object key is deterministic per runtime session under the prefix', () => { + expect(checkpointObjectKey('rt_abc')).toBe('rtsx-checkpoints/rt_abc.tar.gz'); + expect(checkpointObjectKey('rt_abc')).toBe(checkpointObjectKey('rt_abc')); + expect(checkpointObjectKey('rt_xyz')).not.toBe(checkpointObjectKey('rt_abc')); + }); + + test('memory store round-trips bytes and copies defensively', async () => { + const store = new MemoryCheckpointStore(); + const original = Buffer.from('workspace-bytes'); + await store.put('rt_1', original); + + const fetched = await store.get('rt_1'); + expect(fetched?.toString()).toBe('workspace-bytes'); + /* stored copy is independent of the caller's buffer */ + original.fill(0); + expect((await store.get('rt_1'))?.toString()).toBe('workspace-bytes'); + }); + + test('absent checkpoint returns null', async () => { + const store = new MemoryCheckpointStore(); + expect(await store.get('rt_missing')).toBeNull(); + }); + + test('last-writer-wins on the same key', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_1', Buffer.from('v1')); + await store.put('rt_1', Buffer.from('v2')); + expect((await store.get('rt_1'))?.toString()).toBe('v2'); + }); +}); diff --git a/service/src/runtime-session/checkpoint-store.ts b/service/src/runtime-session/checkpoint-store.ts new file mode 100644 index 0000000..faf0336 --- /dev/null +++ b/service/src/runtime-session/checkpoint-store.ts @@ -0,0 +1,72 @@ +import { Client as MinioClient } from 'minio'; +import { env } from '../config'; + +/** + * Durable storage for session workspace checkpoints. Deterministic key per + * runtime session (`.tar.gz`) so a relaunch can + * find the latest checkpoint even if the registry record was lost. Writes are + * serialized by the session lock, so last-writer-wins is the intended + * semantic; object versioning (Phase 4) adds forensic history on top. + */ +export interface CheckpointStore { + put(runtimeSessionId: string, data: Buffer): Promise; + get(runtimeSessionId: string): Promise; +} + +export function checkpointObjectKey(runtimeSessionId: string): string { + return `${env.CHECKPOINT_PREFIX}${runtimeSessionId}.tar.gz`; +} + +/** 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 } : {}), + }); + this.bucket = process.env.CODEAPI_CHECKPOINT_BUCKET ?? process.env.MINIO_BUCKET ?? 'test-bucket'; + } + + async put(runtimeSessionId: string, data: Buffer): Promise { + await this.client.putObject(this.bucket, checkpointObjectKey(runtimeSessionId), data, data.length, { + 'Content-Type': 'application/x-gtar', + }); + } + + async get(runtimeSessionId: string): Promise { + try { + const stream = await this.client.getObject(this.bucket, checkpointObjectKey(runtimeSessionId)); + const chunks: Buffer[] = []; + for await (const chunk of stream) { + chunks.push(chunk as Buffer); + } + return Buffer.concat(chunks); + } catch (error) { + const code = (error as { code?: string })?.code; + if (code === 'NoSuchKey' || code === 'NotFound') return null; + throw error; + } + } +} + +/** In-memory store for bun tests. */ +export class MemoryCheckpointStore implements CheckpointStore { + readonly objects = new Map(); + + async put(runtimeSessionId: string, data: Buffer): Promise { + this.objects.set(runtimeSessionId, Buffer.from(data)); + } + + async get(runtimeSessionId: string): Promise { + const data = this.objects.get(runtimeSessionId); + return data ? Buffer.from(data) : null; + } +} diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts new file mode 100644 index 0000000..385d9ed --- /dev/null +++ b/service/src/runtime-session/checkpoint.ts @@ -0,0 +1,167 @@ +import axios from 'axios'; +import type { LambdaMicrovmClient } from './lambda-client'; +import type { CheckpointStore } from './checkpoint-store'; +import { + acquireRuntimeSessionLock, + readRuntimeSessionRecord, + releaseRuntimeSessionLock, + writeRuntimeSessionRecord, +} from './registry'; +import { checkpointObjectKey } from './checkpoint-store'; +import { microvmCheckpoints, microvmRestores, microvmCheckpointBytes } from '../metrics'; +import logger from '../logger'; + +/** + * 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; +} + +export async function pullCheckpoint( + client: LambdaMicrovmClient, + args: { microvmId: string; endpointBase: string }, + config: CheckpointConfig, +): Promise { + const token = await client.createMicrovmAuthToken({ + microvmId: args.microvmId, + port: config.port, + ttlSeconds: config.authTokenTtlSeconds, + }); + const response = await axios.get(`${args.endpointBase}/api/v2/session/checkpoint`, { + headers: { [token.headerName]: token.token }, + responseType: 'arraybuffer', + maxContentLength: config.maxBytes, + timeout: config.timeoutMs, + }); + return Buffer.from(response.data); +} + +export async function pushRestore( + client: LambdaMicrovmClient, + args: { microvmId: string; endpointBase: string }, + data: Buffer, + config: CheckpointConfig, +): Promise { + const token = await client.createMicrovmAuthToken({ + microvmId: args.microvmId, + port: config.port, + ttlSeconds: config.authTokenTtlSeconds, + }); + await axios.post(`${args.endpointBase}/api/v2/session/restore`, data, { + headers: { + [token.headerName]: token.token, + '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: { + client: LambdaMicrovmClient; + 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 data = await pullCheckpoint(args.client, { + microvmId: record.microvm_id, + endpointBase: args.normalizeEndpoint(record.endpoint), + }, args.config); + await args.store.put(args.runtimeSessionId, data); + microvmCheckpointBytes.observe(data.length); + await writeRuntimeSessionRecord({ + ...record, + workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId), + checkpointed_at: Date.now(), + }, lockToken); + 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: { + client: LambdaMicrovmClient; + store: CheckpointStore; + runtimeSessionId: string; + microvmId: string; + endpointBase: string; + config: CheckpointConfig; +}): Promise<'restored' | 'absent' | 'failed'> { + let data: Buffer | null; + try { + data = await args.store.get(args.runtimeSessionId); + } catch (error) { + 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 'failed'; + } + if (!data) { + microvmRestores.inc({ outcome: 'absent' }); + return 'absent'; + } + try { + await pushRestore(args.client, { microvmId: args.microvmId, endpointBase: args.endpointBase }, data, args.config); + microvmRestores.inc({ outcome: 'restored' }); + logger.info('Session workspace restored from checkpoint', { + runtimeSessionId: args.runtimeSessionId, + bytes: data.length, + }); + return 'restored'; + } catch (error) { + microvmRestores.inc({ outcome: 'failed' }); + logger.warn('Checkpoint restore failed; continuing with a fresh workspace', { + runtimeSessionId: args.runtimeSessionId, + error: error instanceof Error ? error.message : String(error), + }); + return 'failed'; + } +} diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index 7a40a2b..449de96 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -37,6 +37,7 @@ export interface RuntimeSessionRecord { last_seen_at: number; hard_deadline_at?: number; workspace_checkpoint?: string; + checkpointed_at?: number; last_error?: string; } diff --git a/service/src/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts index 6161c1b..0c84230 100644 --- a/service/src/sandbox-backend/index.ts +++ b/service/src/sandbox-backend/index.ts @@ -1,6 +1,7 @@ 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'; @@ -34,7 +35,17 @@ function createBackend(): SandboxBackend { 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(); diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 1fd2939..0fea7e4 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -12,7 +12,9 @@ import { readRuntimeSessionRecord, resetRedisForTests as resetRegistryRedis, setRedisForTests as setRegistryRedis, + writeRuntimeSessionRecord, } from '../runtime-session/registry'; +import { MemoryCheckpointStore, checkpointObjectKey } from '../runtime-session/checkpoint-store'; import { LambdaMicrovmSandboxBackend, normalizeMicrovmEndpoint, type LambdaMicrovmBackendConfig } from './lambda-microvm'; import { SandboxBackendError } from './types'; import type { SandboxExecuteContext, SandboxTransportRequest } from './types'; @@ -25,6 +27,7 @@ let captured: CapturedRequest[] = []; let healthStatus = 200; let executeDelayMs = 0; let mock: InstanceType; +const checkpointBlob = 'FAKE_TAR_GZ_BYTES'; const EXECUTE_RESPONSE = { session_id: 'sess_exec_1', @@ -62,6 +65,15 @@ beforeAll(() => { 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 }); }, }); @@ -101,15 +113,22 @@ function config(overrides: Partial = {}): LambdaMicr 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): LambdaMicrovmSandboxBackend { +function makeBackend( + fake: FakeLambdaMicrovmClient, + cfg?: Partial, + checkpointStore?: MemoryCheckpointStore, +): LambdaMicrovmSandboxBackend { return new LambdaMicrovmSandboxBackend({ clientFactory: () => Promise.resolve(fake), config: config(cfg), pollIntervalMs: 5, + checkpointStore, }); } @@ -359,3 +378,89 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); }); }); + +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); + expect(store.objects.get('rt_ckpt_1')?.toString()).toBe(checkpointBlob); + const record = await readRuntimeSessionRecord('rt_ckpt_1'); + expect(record?.workspace_checkpoint).toBe(checkpointObjectKey('rt_ckpt_1')); + 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', 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'), + }, 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')); + 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 index 857dbe7..3eba9af 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -3,8 +3,11 @@ import { nanoid } from 'nanoid'; import type { LambdaMicrovmClient, 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 } from '../runtime-session/lambda-client'; import { MicrovmOpThrottledError, acquireOpBudget, poisonOpBucket } from '../runtime-session/throttle'; +import { checkpointSession, restoreSession } from '../runtime-session/checkpoint'; import { allocateRuntimeSessionGeneration, readRuntimeSessionRecord, @@ -49,12 +52,18 @@ export interface LambdaMicrovmBackendConfig { 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)); @@ -91,11 +100,17 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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 { @@ -164,8 +179,11 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { * generation, deadline, and image fields. */ const now = Date.now(); const settled = await readRuntimeSessionRecord(runtimeSessionId); - if (settled) { - await writeRuntimeSessionRecord({ ...settled, state: 'RUNNING', last_seen_at: now }, lockToken); + const nextRecord = settled + ? await this.checkpointUnderLock(client, settled, runtimeSessionId, now, lockToken) + : undefined; + if (nextRecord) { + await writeRuntimeSessionRecord(nextRecord, lockToken); } await touchRuntimeSessionActive(runtimeSessionId, now); return result; @@ -237,9 +255,56 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { await this.terminate(client, vm.microvmId, 'error'); throw new SandboxBackendError('MICROVM_FENCED', `Lost session lock for ${runtimeSessionId} after launch`); } + + /* Fresh VM for an existing session: restore its predecessor's workspace + * before the first execute, so an 8h rollover / eviction is invisible. + * A prior record (or a checkpoint pointer) means the session existed. */ + if (this.checkpointStore && this.checkpointsActive() && (record?.workspace_checkpoint || record != null)) { + await restoreSession({ + client, + store: this.checkpointStore, + runtimeSessionId, + microvmId: vm.microvmId, + endpointBase: normalizeMicrovmEndpoint(vm.endpoint ?? ''), + config: this.config.checkpoint, + }); + } return vm; } + /** + * 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({ + client, + store: this.checkpointStore, + runtimeSessionId, + config: this.config.checkpoint, + normalizeEndpoint: normalizeMicrovmEndpoint, + lockToken, + }); + /* checkpointSession wrote the pointer under our lock on success; re-read + * so we don't clobber it with our stale `base`. */ + if (result === 'stored') { + return (await readRuntimeSessionRecord(runtimeSessionId)) ?? base; + } + return base; + } + private async proxyExecute( client: LambdaMicrovmClient, vm: MicrovmDescription, From a71fcfb08a0cd29c27e7771ec9b78385baedd340 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 07:37:42 -0400 Subject: [PATCH 09/24] feat: hookless per-request session binding via X-Runtime-Session-Id header Deliver session mode per /execute request instead of the /run lifecycle hook. Lambda image build hooks only route on the snapshot-compatible base container image, and enabling any runtime hook forces the /ready build hook (which never reaches a stock container listener), so hookless image builds are the reliable path. The runner binds its persistent workspace from the header; the backend stamps it on the proxied execute in session mode and drops runHookPayload from RunMicrovm. Verified on a real hookless MicroVM. --- api/src/api/v2.ts | 12 ++++++- api/src/session-workspace.test.ts | 23 ++++++++++++ api/src/session-workspace.ts | 36 +++++++++++++++++-- scripts/build-lambda-microvm-artifact.sh | 10 ++++-- .../sandbox-backend/lambda-microvm.test.ts | 12 ++++--- service/src/sandbox-backend/lambda-microvm.ts | 32 +++++++++++------ 6 files changed, 102 insertions(+), 23 deletions(-) diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index 4574153..14b4819 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -12,7 +12,12 @@ import { activeSandboxExecutions, recordSandboxExecution } from '../metrics'; import { classifySandboxSafeError } from '../safe-error'; import { withSpan } from '../telemetry'; import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; -import { getBoundSessionWorkspace } from '../session-workspace'; +import { + RUNTIME_SESSION_ID_HEADER, + bindSessionWorkspace, + getBoundSessionWorkspace, + parseSessionBindingFromHeader, +} from '../session-workspace'; import { streamSessionCheckpoint, restoreSessionCheckpoint } from '../session-checkpoint'; const router = express.Router(); @@ -132,6 +137,7 @@ function getJob( egressGrantToken?: string, toolCallSocketEnabled = false, isSynthetic = false, + runtimeSessionHeader?: string | string[], ): Job { const { session_id, language, version, args, stdin, files, @@ -173,6 +179,9 @@ function getJob( validateConstraints(body, rt); + const binding = parseSessionBindingFromHeader(runtimeSessionHeader); + if (binding) bindSessionWorkspace(binding); + return new Job({ session_id: session_id ?? null, runtime: rt, @@ -327,6 +336,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(); diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index 86d09d1..8e4c1a5 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -8,6 +8,7 @@ import { bindSessionWorkspace, getBoundSessionWorkspace, parseSessionBinding, + parseSessionBindingFromHeader, resetSessionWorkspaceStateForTests, unbindSessionWorkspace, } from './session-workspace'; @@ -38,6 +39,28 @@ describe('parseSessionBinding (gating)', () => { 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; diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index 66726ed..21b2282 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -20,11 +20,24 @@ import { * * 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-launch `/run` runHookPayload opting in. When neither is - * active, `getBoundSessionWorkspace()` returns undefined and the runner falls - * back to the untouched fresh-per-job path. + * 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'; + +const RUNTIME_SESSION_ID_PATTERN = /^[A-Za-z0-9._:-]{1,128}$/; + export interface SessionBinding { runtimeSessionId: string; } @@ -53,6 +66,23 @@ export function parseSessionBinding(runHookPayload: string | undefined): Session 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; diff --git a/scripts/build-lambda-microvm-artifact.sh b/scripts/build-lambda-microvm-artifact.sh index c81816e..b6d94c3 100755 --- a/scripts/build-lambda-microvm-artifact.sh +++ b/scripts/build-lambda-microvm-artifact.sh @@ -79,15 +79,19 @@ Next (spike item 2): --code-artifact "uri=$key" \\ --base-image-arn arn:aws:lambda:\${AWS_REGION}:aws:microvm-image:al2023-1 \\ --build-role-arn \\ - --hook-port 8080 \\ + --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. -- /ready + /run/resume/suspend/terminate hooks are served on port 8080 at - /aws/lambda-microvms/runtime/v1/*. +- 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 } diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 0fea7e4..5323822 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -285,7 +285,7 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { }); } - test('launches a session VM with the workspace runHookPayload + idlePolicy and records it', async () => { + test('launches a hookless session VM (idlePolicy, no runHookPayload) and stamps the workspace header on execute', async () => { const fake = fakeClient(); const backend = makeBackend(fake); @@ -298,14 +298,16 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { clientToken?: string; maximumDurationSeconds: number; }; - expect(JSON.parse(runArgs.runHookPayload as string)).toEqual({ - runtime_session_id: 'rt_session_1', - session_workspace: true, - }); + /* 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]); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 3eba9af..6308745 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -29,11 +29,13 @@ import { SandboxBackendError } from './types'; import { Jobs } from '../enum'; import logger from '../logger'; -/** Payload delivered to the MicroVM /run hook to activate the runner's - * persistent session workspace (see api/src/session-workspace.ts). */ -function sessionRunHookPayload(runtimeSessionId: string): string { - return JSON.stringify({ runtime_session_id: runtimeSessionId, session_workspace: true }); -} +/** 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; @@ -78,7 +80,6 @@ export function normalizeMicrovmEndpoint(endpoint: string): string { interface LaunchOptions { clientToken: string; - runHookPayload?: string; idlePolicy?: MicrovmIdlePolicy; maxDurationSeconds: number; } @@ -89,8 +90,9 @@ interface LaunchOptions { * - **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, deliver the /run payload that - * activates the runner's persistent workspace, and reuse it across calls. + * `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. */ @@ -173,7 +175,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { try { const existing = await readRuntimeSessionRecord(runtimeSessionId); const vm = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); - const result = await this.proxyExecute(client, vm, req, ctx); + const result = await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); /* 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. */ @@ -227,7 +229,6 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { const launchedAt = Date.now(); const vm = await this.launch(client, ctx, { clientToken: `sess-${runtimeSessionId}-${generation}`, - runHookPayload: sessionRunHookPayload(runtimeSessionId), idlePolicy: { maxIdleSeconds: this.config.idleSeconds, suspendedSeconds: this.config.suspendedSeconds, @@ -310,6 +311,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { vm: MicrovmDescription, req: SandboxTransportRequest, ctx: SandboxExecuteContext, + runtimeSessionId?: string, ): Promise { const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); const token = await client.createMicrovmAuthToken({ @@ -319,6 +321,14 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { }); 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}`, @@ -332,6 +342,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { headers: { ...injectTraceHeaders(req.headers), [token.headerName]: token.token, + ...sessionHeader, }, signal: ctx.signal, }, @@ -372,7 +383,6 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { ingressConnectorArns: this.config.ingressConnectorArns, egressConnectorArns: this.config.egressConnectorArns, maximumDurationSeconds: opts.maxDurationSeconds, - runHookPayload: opts.runHookPayload, idlePolicy: opts.idlePolicy, clientToken: opts.clientToken, }); From b4d785664ce0cc508de370dd4029234c12550e3b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 11:51:45 -0400 Subject: [PATCH 10/24] docs: Lambda MicroVM stateful sessions runbook + Terraform + image helper Operator guide for the optional AWS Lambda MicroVM backend: the cross-repo picture, from-scratch AWS setup, a full config reference, operating modes, alternative AWS methods (base image, checkpoint store, egress, quota), the PTC replay/blocking distinction, and the hard-won runbook gotchas. - docs/lambda-microvm/terraform: prerequisites module (checkpoint + artifact buckets, build + logging-only execution roles with the sts:TagSession / logs:* trust the build needs, CloudWatch log groups, checkpoint access policy). terraform validate + fmt clean. - service/scripts/create-microvm-image.ts: guaranteed-correct hookless CreateMicrovmImage helper (ALL caps + cgroupv2 off baked in), the one provisioning step Terraform can't own. --- docs/lambda-microvm/README.md | 326 ++++++++++++++++++ docs/lambda-microvm/terraform/.gitignore | 6 + docs/lambda-microvm/terraform/README.md | 44 +++ docs/lambda-microvm/terraform/main.tf | 251 ++++++++++++++ docs/lambda-microvm/terraform/outputs.tf | 47 +++ .../terraform/terraform.tfvars.example | 25 ++ docs/lambda-microvm/terraform/variables.tf | 81 +++++ docs/lambda-microvm/terraform/versions.tf | 14 + service/scripts/create-microvm-image.ts | 113 ++++++ 9 files changed, 907 insertions(+) create mode 100644 docs/lambda-microvm/README.md create mode 100644 docs/lambda-microvm/terraform/.gitignore create mode 100644 docs/lambda-microvm/terraform/README.md create mode 100644 docs/lambda-microvm/terraform/main.tf create mode 100644 docs/lambda-microvm/terraform/outputs.tf create mode 100644 docs/lambda-microvm/terraform/terraform.tfvars.example create mode 100644 docs/lambda-microvm/terraform/variables.tf create mode 100644 docs/lambda-microvm/terraform/versions.tf create mode 100644 service/scripts/create-microvm-image.ts diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md new file mode 100644 index 0000000..152f3fb --- /dev/null +++ b/docs/lambda-microvm/README.md @@ -0,0 +1,326 @@ +# 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 +# → 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)). 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_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_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 for runtime VM stdout to reach CloudWatch. | +| `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_USE_SSL=true`). Prefer granting the +checkpoint policy (Terraform output `checkpoint_access_policy_arn`) to your +CodeAPI task role over minting static keys. + +**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 + docs' `/aws/lambda/microvms/`. +- **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 any live VMs, then delete the image +cd service +AWS_PROFILE=... bun -e 'import { LambdaMicrovmsClient, ListMicrovmsCommand, TerminateMicrovmCommand, DeleteMicrovmImageCommand } from "@aws-sdk/client-lambda-microvms"; const c=new LambdaMicrovmsClient({region:"us-east-1"}); const v=await c.send(new ListMicrovmsCommand({})); for (const m of (v.microvms??[]).filter(x=>!/TERMINAT/.test(x.state))) 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..37be981 --- /dev/null +++ b/docs/lambda-microvm/terraform/main.tf @@ -0,0 +1,251 @@ +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 + bucket = "${var.name_prefix}-artifacts-${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 { + apply_server_side_encryption_by_default { + sse_algorithm = "aws:kms" + } + bucket_key_enabled = true + } +} + +# -------------------------------------------------------------------------- +# 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-${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 { + apply_server_side_encryption_by_default { + sse_algorithm = "aws:kms" + } + bucket_key_enabled = true + } +} + +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, +# not the docs' `/aws/lambda/microvms/...`). 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"] + resources = ["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"] + resources = ["${aws_cloudwatch_log_group.runtime.arn}:*"] + } +} + +resource "aws_iam_role_policy" "execution" { + name = "runtime-logs" + role = aws_iam_role.execution.id + policy = data.aws_iam_policy_document.exec_perms.json +} + +# -------------------------------------------------------------------------- +# IAM 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..7e58306 --- /dev/null +++ b/docs/lambda-microvm/terraform/outputs.tf @@ -0,0 +1,47 @@ +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 = "Attach to your CodeAPI task role for checkpoint S3 access (preferred over the IAM user)." + 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..29fb58d --- /dev/null +++ b/docs/lambda-microvm/terraform/variables.tf @@ -0,0 +1,81 @@ +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 = "" +} + +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). Prefer attaching the checkpoint policy to your CodeAPI task + role instead when running on ECS/EKS; set this false in that case. + 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..882805e --- /dev/null +++ b/docs/lambda-microvm/terraform/versions.tf @@ -0,0 +1,14 @@ +terraform { + required_version = ">= 1.5" + + required_providers { + aws = { + source = "hashicorp/aws" + version = ">= 5.40" + } + } +} + +provider "aws" { + region = var.region +} diff --git a/service/scripts/create-microvm-image.ts b/service/scripts/create-microvm-image.ts new file mode 100644 index 0000000..6baf363 --- /dev/null +++ b/service/scripts/create-microvm-image.ts @@ -0,0 +1,113 @@ +/** + * 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); +} + +/* 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' }, +}; + +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(); + 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); + 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); +}); From 3c7e9d380ed27318cf703352affe1f5672100478 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 12:25:44 -0400 Subject: [PATCH 11/24] fix: address review findings (lock TTL, symlink chown, fence ordering, +) Fixes from the Macroscope review pass: - registry: derive RUNTIME_SESSION_LOCK_TTL_MS from the actual launch/health/ execute/checkpoint budgets (was a 60s placeholder, could expire mid-work at defaults); guard readRuntimeSessionRecord JSON.parse so a corrupt key reads as missing instead of wedging the session. - checkpoint: fence (fenced record write) BEFORE the deterministic-key S3 put so a lock-expired caller can't clobber a newer blob; cap restore size against maxBytes before buffering. - session-checkpoint: lchown + never follow symlinks when chowning a restored (untrusted) checkpoint, so a symlinked entry can't re-own files outside the workspace. - lambda-client: throw when a MicroVM response omits microvmId instead of returning '' (a partial RunMicrovm response would otherwise orphan a billable VM behind getMicrovm('')/terminateMicrovm('')). - router: skip runtime_session_hint validation in stateless mode (the field is ignored there, so a malformed hint must not 400). - v2/session: only run in the persistent workspace when THIS request carried a valid X-Runtime-Session-Id, so a headerless request never inherits a prior session's workspace/UID. - entrypoint: only ever raise RLIMIT_NOFILE (guard so it never clamps a higher host default down to 65536). - secure-startup: warn (don't silently no-op) on affinity+http. Deferred with reasons in the PR: zset-orphan (no consumer until the sweeper PR), bindSessionWorkspace rebind race (prevented by the 1-VM-1-session invariant), releaseLock swallow (TTL self-heals), /run applyRunHook (inert in the hookless design), startupApiOnly policy placement (split-deploy design Q). --- api/src/api/v2.ts | 10 ++++--- api/src/entrypoint.sh | 18 +++++++++---- api/src/session-checkpoint.ts | 10 ++++++- service/src/runtime-session/checkpoint.ts | 26 +++++++++++++++--- .../src/runtime-session/lambda-client-aws.ts | 9 ++++++- service/src/runtime-session/registry.test.ts | 5 ++++ service/src/runtime-session/registry.ts | 27 ++++++++++++++----- service/src/secure-startup.ts | 10 +++++++ service/src/service/router.ts | 8 +++++- 9 files changed, 103 insertions(+), 20 deletions(-) diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index 14b4819..eb8c8f2 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -15,7 +15,6 @@ import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; import { RUNTIME_SESSION_ID_HEADER, bindSessionWorkspace, - getBoundSessionWorkspace, parseSessionBindingFromHeader, } from '../session-workspace'; import { streamSessionCheckpoint, restoreSessionCheckpoint } from '../session-checkpoint'; @@ -179,8 +178,13 @@ 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); - if (binding) bindSessionWorkspace(binding); + const session = binding ? bindSessionWorkspace(binding) ?? null : null; return new Job({ session_id: session_id ?? null, @@ -205,7 +209,7 @@ function getJob( egress_grant: egressGrantToken, tool_call_socket_enabled: toolCallSocketEnabled, is_synthetic: isSynthetic, - session: getBoundSessionWorkspace() ?? null, + session, }); } diff --git a/api/src/entrypoint.sh b/api/src/entrypoint.sh index c7f576d..bfdd8c6 100755 --- a/api/src/entrypoint.sh +++ b/api/src/entrypoint.sh @@ -3,14 +3,22 @@ set -e echo "Starting NsJail sandbox API..." -# Raise the RLIMIT_NOFILE hard cap so per-job nsjail can set its own soft +# 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. No-op where the limit is already higher -# (e.g. Docker's ~1M default) or where CAP_SYS_RESOURCE is unavailable. -ulimit -Hn 65536 2>/dev/null || true -ulimit -Sn 65536 2>/dev/null || true +# 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}" diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index c129337..8137044 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -85,10 +85,18 @@ export async function restoreSessionCheckpoint(req: Request, res: Response): Pro } async function chownRecursive(dir: string, uid: number, gid: number): Promise { - await fsp.chown(dir, uid, gid).catch(() => {}); + 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/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts index 385d9ed..4da804b 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -102,13 +102,21 @@ export async function checkpointSession(args: { microvmId: record.microvm_id, endpointBase: args.normalizeEndpoint(record.endpoint), }, args.config); - await args.store.put(args.runtimeSessionId, data); - microvmCheckpointBytes.observe(data.length); - await writeRuntimeSessionRecord({ + /* Fence BEFORE writing the object store. The checkpoint key is + * deterministic per session (last-writer-wins), so a caller whose lock + * expired must not clobber a newer blob. The fenced record write is the + * ownership check: if it reports we were fenced, skip the store entirely. */ + const persisted = await writeRuntimeSessionRecord({ ...record, workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId), checkpointed_at: Date.now(), }, lockToken); + if (!persisted) { + microvmCheckpoints.inc({ outcome: 'skipped_busy' }); + return 'skipped_busy'; + } + await args.store.put(args.runtimeSessionId, data); + microvmCheckpointBytes.observe(data.length); microvmCheckpoints.inc({ outcome: 'stored' }); return 'stored'; } catch (error) { @@ -148,6 +156,18 @@ export async function restoreSession(args: { microvmRestores.inc({ outcome: 'absent' }); return 'absent'; } + /* Guard against an oversized/stray checkpoint before buffering it into the + * restore request. The put side caps size, but a checkpoint written under a + * looser prior config (or externally) could otherwise exhaust the process. */ + if (data.length > args.config.maxBytes) { + microvmRestores.inc({ outcome: 'failed' }); + logger.warn('Checkpoint exceeds maxBytes; continuing with a fresh workspace', { + runtimeSessionId: args.runtimeSessionId, + bytes: data.length, + maxBytes: args.config.maxBytes, + }); + return 'failed'; + } try { await pushRestore(args.client, { microvmId: args.microvmId, endpointBase: args.endpointBase }, data, args.config); microvmRestores.inc({ outcome: 'restored' }); diff --git a/service/src/runtime-session/lambda-client-aws.ts b/service/src/runtime-session/lambda-client-aws.ts index ffb0a17..7798d90 100644 --- a/service/src/runtime-session/lambda-client-aws.ts +++ b/service/src/runtime-session/lambda-client-aws.ts @@ -46,8 +46,15 @@ function toDescription(response: { 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 ?? '', + microvmId: response.microvmId, state: (response.state ?? 'PENDING') as MicrovmLifecycleState, endpoint: response.endpoint, imageArn: response.imageArn, diff --git a/service/src/runtime-session/registry.test.ts b/service/src/runtime-session/registry.test.ts index 8955813..cf30f26 100644 --- a/service/src/runtime-session/registry.test.ts +++ b/service/src/runtime-session/registry.test.ts @@ -84,6 +84,11 @@ describe('fenced record writes', () => { 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); diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index 449de96..bd2608a 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -46,11 +46,18 @@ const LOCK_PREFIX = 'rtsx:lock:'; const GEN_PREFIX = 'rtsx:gen:'; const ACTIVE_ZSET = 'rtsx:active'; -/** Launch budget placeholder until the Lambda backend config lands; the lock - * must outlive lock-wait + launch + execute so a live holder is never fenced - * mid-operation. */ -const LAUNCH_BUDGET_MS = 60_000; -export const RUNTIME_SESSION_LOCK_TTL_MS = env.JOB_TIMEOUT + LAUNCH_BUDGET_MS + 60_000; +/** The session lock is held across the whole `executeSession` critical path — + * launch + health + execute + post-run checkpoint — so its TTL must outlive the + * sum of those configurable budgets, else a live holder is fenced mid-operation + * and a second worker can acquire the same session concurrently. Derive it from + * the actual budgets (not a fixed placeholder) so raising any one of them keeps + * the lock safe, plus headroom for lock-wait + scheduling jitter. */ +export const RUNTIME_SESSION_LOCK_TTL_MS = + env.JOB_TIMEOUT + + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + + env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS + + 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; @@ -151,7 +158,15 @@ export async function releaseRuntimeSessionLock(runtimeSessionId: string, token: export async function readRuntimeSessionRecord(runtimeSessionId: string): Promise { const data = await redis.get(`${SESS_PREFIX}${runtimeSessionId}`); - return data != null ? (JSON.parse(data) as RuntimeSessionRecord) : null; + 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 diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 218cb7e..0bcd1b2 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 { @@ -58,6 +59,15 @@ export function validateSandboxBackendPolicy(): void { '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') { diff --git a/service/src/service/router.ts b/service/src/service/router.ts index aa57646..f914662 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -136,11 +136,17 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) 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: validateRuntimeSessionHint(body.runtime_session_hint), + hint, }); } catch (error) { if (error instanceof RuntimeSessionHintError) { From 02cbdf1b6de82a5b4598ef3485d71c2239c1f396 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 12:37:15 -0400 Subject: [PATCH 12/24] fix(docs): address review findings in Lambda MicroVM IaC + guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - terraform: artifact + checkpoint buckets use SSE-S3 (AES256) so the s3:GetObject build role and the checkpoint access policy need no kms:Decrypt grant (SSE-KMS would AccessDenied on read). - terraform: build + runtime log policies grant both the log-group ARN and its stream form — stream-level actions (CreateLogStream/PutLogEvents) don't match the group ARN, which would fail builds with an empty stateReason. - terraform: validate artifact_bucket_name is non-empty when reusing an existing bucket (else the policy resolves to arn:aws:s3:::/*); bump required_version to >= 1.9 for cross-variable validation. - create-microvm-image.ts: cap the poll loop with a deadline (default 30m, MICROVM_BUILD_DEADLINE_MINUTES) and exit non-zero, so a wedged CREATING build can't hang a provisioning job forever. - README: add MINIO_PORT=443 to the S3 example (client defaults to 9000); scope the teardown sweep to VMs from this image's ARN so it can't terminate unrelated MicroVMs in a shared account. --- docs/lambda-microvm/README.md | 8 +++- docs/lambda-microvm/terraform/main.tf | 43 ++++++++++++++++------ docs/lambda-microvm/terraform/variables.tf | 7 ++++ docs/lambda-microvm/terraform/versions.tf | 3 +- service/scripts/create-microvm-image.ts | 8 ++++ 5 files changed, 54 insertions(+), 15 deletions(-) diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md index 152f3fb..eefec6c 100644 --- a/docs/lambda-microvm/README.md +++ b/docs/lambda-microvm/README.md @@ -145,6 +145,7 @@ LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS=arn:aws:lambda:us-east-1:aws:network-connec # 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 @@ -313,9 +314,12 @@ Each of these cost a silent or blind failure during bring-up: ## Teardown ```bash -# terminate any live VMs, then delete the image +# 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 -AWS_PROFILE=... bun -e 'import { LambdaMicrovmsClient, ListMicrovmsCommand, TerminateMicrovmCommand, DeleteMicrovmImageCommand } from "@aws-sdk/client-lambda-microvms"; const c=new LambdaMicrovmsClient({region:"us-east-1"}); const v=await c.send(new ListMicrovmsCommand({})); for (const m of (v.microvms??[]).filter(x=>!/TERMINAT/.test(x.state))) await c.send(new TerminateMicrovmCommand({microvmIdentifier:m.microvmId})); await c.send(new DeleteMicrovmImageCommand({imageIdentifier:"codeapi-session"})).catch(()=>{});' +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 diff --git a/docs/lambda-microvm/terraform/main.tf b/docs/lambda-microvm/terraform/main.tf index 37be981..ac7d519 100644 --- a/docs/lambda-microvm/terraform/main.tf +++ b/docs/lambda-microvm/terraform/main.tf @@ -33,10 +33,14 @@ 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 = "aws:kms" + sse_algorithm = "AES256" } - bucket_key_enabled = true } } @@ -70,10 +74,13 @@ resource "aws_s3_bucket_versioning" "checkpoint" { 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 = "aws:kms" + sse_algorithm = "AES256" } - bucket_key_enabled = true } } @@ -143,10 +150,16 @@ data "aws_iam_policy_document" "build_perms" { } statement { - sid = "BuildLogs" - effect = "Allow" - actions = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"] - resources = ["arn:aws:logs:${var.region}:${local.account_id}:log-group:/aws/lambda-microvms/*"] + 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" { @@ -195,10 +208,16 @@ resource "aws_iam_role" "execution" { data "aws_iam_policy_document" "exec_perms" { statement { - sid = "RuntimeLogs" - effect = "Allow" - actions = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"] - resources = ["${aws_cloudwatch_log_group.runtime.arn}:*"] + 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}:*", + ] } } diff --git a/docs/lambda-microvm/terraform/variables.tf b/docs/lambda-microvm/terraform/variables.tf index 29fb58d..f10af79 100644 --- a/docs/lambda-microvm/terraform/variables.tf +++ b/docs/lambda-microvm/terraform/variables.tf @@ -35,6 +35,13 @@ 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" { diff --git a/docs/lambda-microvm/terraform/versions.tf b/docs/lambda-microvm/terraform/versions.tf index 882805e..26471d8 100644 --- a/docs/lambda-microvm/terraform/versions.tf +++ b/docs/lambda-microvm/terraform/versions.tf @@ -1,5 +1,6 @@ terraform { - required_version = ">= 1.5" + # >= 1.9 for cross-variable references in variable validation blocks. + required_version = ">= 1.9" required_providers { aws = { diff --git a/service/scripts/create-microvm-image.ts b/service/scripts/create-microvm-image.ts index 6baf363..4c52409 100644 --- a/service/scripts/create-microvm-image.ts +++ b/service/scripts/create-microvm-image.ts @@ -83,12 +83,20 @@ async function main(): Promise { } 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`); From 8e2da55797b15ffc7bd9f4940ab1fa2996ad2d1d Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 12:46:57 -0400 Subject: [PATCH 13/24] fix(docs): region-unique bucket names + correct checkpoint-cred guidance - terraform: include region in the artifact + checkpoint bucket names (S3 names are global, so a same-prefix second-region apply would collide). - docs: correct the checkpoint credential guidance. The MinIO-compatible client reads only static MINIO_ACCESS_KEY/SECRET and does not load task-role/IRSA creds, so attaching checkpoint_access_policy_arn to a task role alone does not work. Point operators at create_checkpoint_access_user (static keys) and note IRSA-aware loading as a follow-up. Fixed across variables.tf, outputs.tf, README. --- docs/lambda-microvm/README.md | 9 ++++++--- docs/lambda-microvm/terraform/main.tf | 8 +++++--- docs/lambda-microvm/terraform/outputs.tf | 8 +++++++- docs/lambda-microvm/terraform/variables.tf | 6 ++++-- 4 files changed, 22 insertions(+), 9 deletions(-) diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md index eefec6c..6f62e3e 100644 --- a/docs/lambda-microvm/README.md +++ b/docs/lambda-microvm/README.md @@ -249,9 +249,12 @@ 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_USE_SSL=true`). Prefer granting the -checkpoint policy (Terraform output `checkpoint_access_policy_arn`) to your -CodeAPI task role over minting static keys. +`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 diff --git a/docs/lambda-microvm/terraform/main.tf b/docs/lambda-microvm/terraform/main.tf index ac7d519..5e45fd9 100644 --- a/docs/lambda-microvm/terraform/main.tf +++ b/docs/lambda-microvm/terraform/main.tf @@ -14,8 +14,10 @@ locals { # S3: code-artifact bucket (the zip that create-microvm-image reads) # -------------------------------------------------------------------------- resource "aws_s3_bucket" "artifact" { - count = var.create_artifact_bucket ? 1 : 0 - bucket = "${var.name_prefix}-artifacts-${local.account_id}" + 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 } @@ -51,7 +53,7 @@ resource "aws_s3_bucket_server_side_encryption_configuration" "artifact" { # resumable cache rather than a system of record. # -------------------------------------------------------------------------- resource "aws_s3_bucket" "checkpoint" { - bucket = "${var.name_prefix}-checkpoints-${local.account_id}" + bucket = "${var.name_prefix}-checkpoints-${var.region}-${local.account_id}" force_destroy = true tags = local.base_tags } diff --git a/docs/lambda-microvm/terraform/outputs.tf b/docs/lambda-microvm/terraform/outputs.tf index 7e58306..3f6dfee 100644 --- a/docs/lambda-microvm/terraform/outputs.tf +++ b/docs/lambda-microvm/terraform/outputs.tf @@ -31,7 +31,13 @@ output "runtime_log_group" { } output "checkpoint_access_policy_arn" { - description = "Attach to your CodeAPI task role for checkpoint S3 access (preferred over the IAM user)." + 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 } diff --git a/docs/lambda-microvm/terraform/variables.tf b/docs/lambda-microvm/terraform/variables.tf index f10af79..4258e96 100644 --- a/docs/lambda-microvm/terraform/variables.tf +++ b/docs/lambda-microvm/terraform/variables.tf @@ -74,8 +74,10 @@ variable "create_checkpoint_access_user" { description = <<-EOT Create an IAM user + access key with read/write on the checkpoint bucket, for the CodeAPI service's MinIO-compatible checkpoint client (MINIO_ACCESS_KEY / - MINIO_SECRET_KEY). Prefer attaching the checkpoint policy to your CodeAPI task - role instead when running on ECS/EKS; set this false in that case. + 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 From daa43b06e596fc300bf846e881116177283f5de7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 13:21:49 -0400 Subject: [PATCH 14/24] fix: address Codex review of stateful session mechanics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Bind the session on checkpoint/restore (F1): the hookless path runs /session/restore before the first /execute, so the runner had nothing bound and every real restore 409'd, losing checkpoint state across VM expiry. The runner now binds from X-Runtime-Session-Id in the checkpoint/restore routes, and the backend sends that header on both proxied calls. - Clear/terminate a session VM on reuse failure or abort (F2/F5): on a dead reused VM (idlePolicy auto-terminated) or an aborted execute (runner keeps NsJail alive after the socket closes), terminate the VM and drop the record so the next call relaunches + restores instead of reusing a dead-or-dirty VM. A plain non-200 from a live runner leaves the warm VM intact. - Enforce checkpoint size before buffering (F7): CheckpointStore.get takes maxBytes and stats the S3 object first, so an oversized/stray checkpoint can't OOM the worker before the (now-removed, too-late) post-buffer guard. - Mark session outputs surfaced only after upload (F3): a dropped upload no longer permanently suppresses an unchanged file next turn. - Bake runner file/egress/manifest env into the image (F4): create-microvm-image takes --env-json so images can reach FILE_SERVER_URL / EGRESS_GATEWAY_URL instead of building invalid /sessions/... URLs. Deferred (P2, design choice): don't count checkpoint time in the client job timeout — decoupling checkpoint from the response path needs a call on where it sits relative to job completion; raised with the maintainer. --- api/src/api/v2.ts | 19 +++++++++- api/src/job.ts | 19 +++++++++- docs/lambda-microvm/README.md | 12 +++++- service/scripts/create-microvm-image.ts | 17 ++++++++- .../runtime-session/checkpoint-store.test.ts | 18 ++++++--- .../src/runtime-session/checkpoint-store.ts | 37 ++++++++++++++----- service/src/runtime-session/checkpoint.ts | 35 +++++++++--------- .../sandbox-backend/lambda-microvm.test.ts | 5 +++ service/src/sandbox-backend/lambda-microvm.ts | 36 +++++++++++++++++- 9 files changed, 159 insertions(+), 39 deletions(-) diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index eb8c8f2..649e766 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -465,7 +465,22 @@ router.get('/runtimes', (_req: Request, res: Response) => { * 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. */ -router.get('/session/checkpoint', (_req: Request, res: Response) => streamSessionCheckpoint(res)); -router.post('/session/restore', (req: Request, res: Response) => restoreSessionCheckpoint(req, res)); +/* 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). */ +function bindSessionFromHeader(req: Request): void { + const binding = parseSessionBindingFromHeader(req.headers[RUNTIME_SESSION_ID_HEADER]); + if (binding) bindSessionWorkspace(binding); +} + +router.get('/session/checkpoint', (req: Request, res: Response) => { + bindSessionFromHeader(req); + return streamSessionCheckpoint(res); +}); +router.post('/session/restore', (req: Request, res: Response) => { + bindSessionFromHeader(req); + return restoreSessionCheckpoint(req, res); +}); export default router; diff --git a/api/src/job.ts b/api/src/job.ts index 3839b44..e40c424 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -676,6 +676,11 @@ 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(); @@ -1588,7 +1593,10 @@ export class Job { } this.sessionFiles.push(fileData); this.generatedFiles.push({ id: newId, name: relativePath, path: fullPath }); - if (this.session) this.session.markSurfaced(relativePath, outputSignature); + /* 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) this.pendingSurfaced.set(newId, { name: relativePath, signature: outputSignature }); return { collected: true, truncated: false, stopLoop: false }; } @@ -1755,6 +1763,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 }, diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md index 6f62e3e..4818743 100644 --- a/docs/lambda-microvm/README.md +++ b/docs/lambda-microvm/README.md @@ -122,13 +122,21 @@ 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 + --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)). To ship new runner code later, re-run +[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 diff --git a/service/scripts/create-microvm-image.ts b/service/scripts/create-microvm-image.ts index 4c52409..bef3bbb 100644 --- a/service/scripts/create-microvm-image.ts +++ b/service/scripts/create-microvm-image.ts @@ -56,6 +56,21 @@ if (!artifactUri || !buildRoleArn) { 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. @@ -69,7 +84,7 @@ const shared = { cpuConfigurations: [{ architecture: 'ARM_64' as const }], resources: [{ minimumMemoryInMiB: memory }], additionalOsCapabilities: ['ALL' as const], - environmentVariables: { SANDBOX_USE_CGROUPV2: 'false' }, + environmentVariables: { SANDBOX_USE_CGROUPV2: 'false', ...parseEnvJson() }, }; const client = new LambdaMicrovmsClient({ region, retryMode: 'adaptive', maxAttempts: 3 }); diff --git a/service/src/runtime-session/checkpoint-store.test.ts b/service/src/runtime-session/checkpoint-store.test.ts index cd261fc..635945e 100644 --- a/service/src/runtime-session/checkpoint-store.test.ts +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -1,5 +1,7 @@ import { describe, expect, test } from 'bun:test'; -import { MemoryCheckpointStore, checkpointObjectKey } from './checkpoint-store'; +import { MemoryCheckpointStore, CheckpointTooLargeError, checkpointObjectKey } from './checkpoint-store'; + +const BIG = 1_000_000; describe('checkpoint store', () => { test('object key is deterministic per runtime session under the prefix', () => { @@ -13,22 +15,28 @@ describe('checkpoint store', () => { const original = Buffer.from('workspace-bytes'); await store.put('rt_1', original); - const fetched = await store.get('rt_1'); + 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'))?.toString()).toBe('workspace-bytes'); + 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')).toBeNull(); + expect(await store.get('rt_missing', BIG)).toBeNull(); }); test('last-writer-wins on the same key', async () => { const store = new MemoryCheckpointStore(); await store.put('rt_1', Buffer.from('v1')); await store.put('rt_1', Buffer.from('v2')); - expect((await store.get('rt_1'))?.toString()).toBe('v2'); + expect((await store.get('rt_1', BIG))?.toString()).toBe('v2'); + }); + + test('rejects a checkpoint larger than maxBytes', async () => { + const store = new MemoryCheckpointStore(); + await store.put('rt_big', 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 index faf0336..71add1a 100644 --- a/service/src/runtime-session/checkpoint-store.ts +++ b/service/src/runtime-session/checkpoint-store.ts @@ -10,9 +10,14 @@ import { env } from '../config'; */ export interface CheckpointStore { put(runtimeSessionId: string, data: Buffer): Promise; - get(runtimeSessionId: string): Promise; + /** `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; } +export class CheckpointTooLargeError extends Error {} + export function checkpointObjectKey(runtimeSessionId: string): string { return `${env.CHECKPOINT_PREFIX}${runtimeSessionId}.tar.gz`; } @@ -41,19 +46,27 @@ export class MinioCheckpointStore implements CheckpointStore { }); } - async get(runtimeSessionId: string): Promise { + async get(runtimeSessionId: string, maxBytes: number): Promise { + const key = checkpointObjectKey(runtimeSessionId); + let size: number; try { - const stream = await this.client.getObject(this.bucket, checkpointObjectKey(runtimeSessionId)); - const chunks: Buffer[] = []; - for await (const chunk of stream) { - chunks.push(chunk as Buffer); - } - return Buffer.concat(chunks); + 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[] = []; + for await (const chunk of stream) { + chunks.push(chunk as Buffer); + } + return Buffer.concat(chunks); } } @@ -65,8 +78,12 @@ export class MemoryCheckpointStore implements CheckpointStore { this.objects.set(runtimeSessionId, Buffer.from(data)); } - async get(runtimeSessionId: string): Promise { + async get(runtimeSessionId: string, maxBytes: number): Promise { const data = this.objects.get(runtimeSessionId); - return data ? Buffer.from(data) : null; + if (!data) return null; + 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 index 4da804b..e11e9a9 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -30,9 +30,15 @@ export interface CheckpointConfig { 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( client: LambdaMicrovmClient, - args: { microvmId: string; endpointBase: string }, + args: { microvmId: string; endpointBase: string; runtimeSessionId: string }, config: CheckpointConfig, ): Promise { const token = await client.createMicrovmAuthToken({ @@ -41,7 +47,10 @@ export async function pullCheckpoint( ttlSeconds: config.authTokenTtlSeconds, }); const response = await axios.get(`${args.endpointBase}/api/v2/session/checkpoint`, { - headers: { [token.headerName]: token.token }, + headers: { + [token.headerName]: token.token, + [RUNTIME_SESSION_ID_HEADER]: args.runtimeSessionId, + }, responseType: 'arraybuffer', maxContentLength: config.maxBytes, timeout: config.timeoutMs, @@ -51,7 +60,7 @@ export async function pullCheckpoint( export async function pushRestore( client: LambdaMicrovmClient, - args: { microvmId: string; endpointBase: string }, + args: { microvmId: string; endpointBase: string; runtimeSessionId: string }, data: Buffer, config: CheckpointConfig, ): Promise { @@ -63,6 +72,7 @@ export async function pushRestore( await axios.post(`${args.endpointBase}/api/v2/session/restore`, data, { headers: { [token.headerName]: token.token, + [RUNTIME_SESSION_ID_HEADER]: args.runtimeSessionId, 'Content-Type': 'application/x-gtar', }, maxBodyLength: config.maxBytes, @@ -101,6 +111,7 @@ export async function checkpointSession(args: { const data = await pullCheckpoint(args.client, { microvmId: record.microvm_id, endpointBase: args.normalizeEndpoint(record.endpoint), + runtimeSessionId: args.runtimeSessionId, }, args.config); /* Fence BEFORE writing the object store. The checkpoint key is * deterministic per session (last-writer-wins), so a caller whose lock @@ -141,9 +152,11 @@ export async function restoreSession(args: { endpointBase: string; config: CheckpointConfig; }): Promise<'restored' | 'absent' | '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. */ let data: Buffer | null; try { - data = await args.store.get(args.runtimeSessionId); + data = await args.store.get(args.runtimeSessionId, args.config.maxBytes); } catch (error) { microvmRestores.inc({ outcome: 'failed' }); logger.warn('Checkpoint fetch failed; continuing with a fresh workspace', { @@ -156,20 +169,8 @@ export async function restoreSession(args: { microvmRestores.inc({ outcome: 'absent' }); return 'absent'; } - /* Guard against an oversized/stray checkpoint before buffering it into the - * restore request. The put side caps size, but a checkpoint written under a - * looser prior config (or externally) could otherwise exhaust the process. */ - if (data.length > args.config.maxBytes) { - microvmRestores.inc({ outcome: 'failed' }); - logger.warn('Checkpoint exceeds maxBytes; continuing with a fresh workspace', { - runtimeSessionId: args.runtimeSessionId, - bytes: data.length, - maxBytes: args.config.maxBytes, - }); - return 'failed'; - } try { - await pushRestore(args.client, { microvmId: args.microvmId, endpointBase: args.endpointBase }, data, args.config); + await pushRestore(args.client, { microvmId: 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, diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 5323822..01b1888 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -402,6 +402,9 @@ describe('LambdaMicrovmSandboxBackend auto-checkpoint', () => { 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(store.objects.get('rt_ckpt_1')?.toString()).toBe(checkpointBlob); const record = await readRuntimeSessionRecord('rt_ckpt_1'); expect(record?.workspace_checkpoint).toBe(checkpointObjectKey('rt_ckpt_1')); @@ -429,6 +432,8 @@ describe('LambdaMicrovmSandboxBackend auto-checkpoint', () => { /* 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); }); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 6308745..9695819 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -12,6 +12,7 @@ import { allocateRuntimeSessionGeneration, readRuntimeSessionRecord, releaseRuntimeSessionLock, + removeRuntimeSession, touchRuntimeSessionActive, waitForRuntimeSessionLock, writeRuntimeSessionRecord, @@ -175,7 +176,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { try { const existing = await readRuntimeSessionRecord(runtimeSessionId); const vm = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); - const result = await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); + const result = await this.executeOnSessionVm(client, vm, req, ctx, runtimeSessionId, lockToken); /* 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. */ @@ -194,6 +195,39 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { } } + /** + * 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, + ): Promise { + try { + return await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); + } catch (error) { + const sandboxResponded = error instanceof Error && error.message === 'Error from sandbox'; + if (ctx.signal.aborted || !sandboxResponded) { + await this.terminate(client, vm.microvmId, ctx.signal.aborted ? 'timeout' : 'error').catch(() => {}); + await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); + } + throw error; + } + } + private async findOrLaunchSession( client: LambdaMicrovmClient, ctx: SandboxExecuteContext, From 59f6faeb83869309282246b8f0dd6eaa88e09295 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 13:52:56 -0400 Subject: [PATCH 15/24] fix: address Macroscope review of the Codex-fix commit - executeOnSessionVm error classification (High, regression from the prior commit): a healthy runner returning a non-2xx makes axios throw an AxiosError carrying `.response`, which the old `message === 'Error from sandbox'` check treated as unreachable and tore the VM down. Now keep the VM when the runner responded (isAxiosError && response), terminate only on no-response (connection/timeout) or abort. Regression test added. - findOrLaunchSession: treat an image_arn/version/port mismatch as non-reusable so a warm session relaunches on the current config after a deploy instead of running the old image (or health-checking the wrong port -> UNHEALTHY). - MinioCheckpointStore.get: cap accumulated bytes during download too, not just via statObject, so an object that grows between stat and read can't OOM. - /session/checkpoint + /session/restore: fail closed (409) when the request carries no valid X-Runtime-Session-Id instead of acting on a stale bound session, and `.catch(next)` so a rejected handler promise is forwarded to Express 4's error handler instead of hanging the request. --- api/src/api/v2.ts | 29 +++++++++++++------ .../src/runtime-session/checkpoint-store.ts | 7 +++++ .../sandbox-backend/lambda-microvm.test.ts | 18 +++++++++++- service/src/sandbox-backend/lambda-microvm.ts | 19 +++++++++++- 4 files changed, 62 insertions(+), 11 deletions(-) diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index 649e766..53b7883 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -468,19 +468,30 @@ router.get('/runtimes', (_req: Request, res: Response) => { /* 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). */ -function bindSessionFromHeader(req: Request): void { + * 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) bindSessionWorkspace(binding); + if (!binding) return false; + bindSessionWorkspace(binding); + return true; } -router.get('/session/checkpoint', (req: Request, res: Response) => { - bindSessionFromHeader(req); - return streamSessionCheckpoint(res); +/* 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) => { - bindSessionFromHeader(req); - return restoreSessionCheckpoint(req, res); +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/service/src/runtime-session/checkpoint-store.ts b/service/src/runtime-session/checkpoint-store.ts index 71add1a..f09fd47 100644 --- a/service/src/runtime-session/checkpoint-store.ts +++ b/service/src/runtime-session/checkpoint-store.ts @@ -63,7 +63,14 @@ export class MinioCheckpointStore implements CheckpointStore { } 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); diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 01b1888..5038b0c 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -26,6 +26,7 @@ let server: ReturnType; let captured: CapturedRequest[] = []; let healthStatus = 200; let executeDelayMs = 0; +let executeStatus = 200; let mock: InstanceType; const checkpointBlob = 'FAKE_TAR_GZ_BYTES'; @@ -61,7 +62,7 @@ beforeAll(() => { await new Promise((resolve) => setTimeout(resolve, executeDelayMs)); } return new Response(JSON.stringify(EXECUTE_RESPONSE), { - status: 200, + status: executeStatus, headers: { 'Content-Type': 'application/json' }, }); } @@ -92,6 +93,7 @@ beforeEach(async () => { captured = []; healthStatus = 200; executeDelayMs = 0; + executeStatus = 200; }); afterEach(() => { @@ -327,6 +329,20 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { expect(executes).toHaveLength(2); }); + 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('two concurrent executions on one session serialize on the registry lock', async () => { const fake = fakeClient(); const backend = makeBackend(fake); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 9695819..c6caf51 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -219,7 +219,15 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { try { return await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); } catch (error) { - const sandboxResponded = error instanceof Error && error.message === 'Error from sandbox'; + /* Keep the warm VM only when the runner actually responded — a non-2xx + * makes axios throw an AxiosError carrying `.response`, so the VM is alive + * and just the request failed. No response (connection/timeout/abort) or a + * failed health check means the VM is unreachable/dirty: terminate it and + * drop the record so the next call relaunches + restores. (`Error from + * sandbox` covers proxyExecute's manual 2xx-but-not-200 throw.) */ + const sandboxResponded = + (axios.isAxiosError(error) && error.response != null) || + (error instanceof Error && error.message === 'Error from sandbox'); if (ctx.signal.aborted || !sandboxResponded) { await this.terminate(client, vm.microvmId, ctx.signal.aborted ? 'timeout' : 'error').catch(() => {}); await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); @@ -236,10 +244,19 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { lockToken: string, ): Promise { 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; const reusable = record && record.state === 'RUNNING' && record.microvm_id && record.endpoint + && configMatches && (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 From e0dd9121b21febf988b9ae7a0a25c24fdcee5673 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 14:25:58 -0400 Subject: [PATCH 16/24] fix: address Codex review round 2 of stateful session mechanics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Lock TTL (P1): budget 2x launch (throttle wait + poll-to-RUNNING) and 2x checkpoint (restore-before-exec + post-exec) so a slow relaunch+restore can't expire the lock mid-critical-section. - Read-only inputs re-download (P1): a read_only prime (skill file) is reported as not-primed so it always re-downloads, restoring pristine content + the 0444/root protection a sandbox could have stripped by unlink+replace. - Relaunch idle-expired sessions (P1): treat a record whose last_seen is past idle+suspended as non-reusable, so the first request after idle expiry relaunches+restores instead of reusing a dead endpoint and 503ing. - Content-hash output diffing (P2): suppress a surfaced output only when its CONTENT hash is unchanged, not size+mtime (spoofable via os.utime); the hash is computed once per file and reused for input-modification detection. - RunMicrovm logging (P2): send the CloudWatch logging config (new LAMBDA_MICROVM_LOG_GROUP) so VM stdout reaches CloudWatch — the role alone wasn't enough. - Preserve sessions on non-transport failures (P2): terminate the VM only on abort or a transport-level axios failure (no response); a throttled CreateMicrovmAuthToken or a non-2xx runner response keeps the warm VM. - Skip the optional checkpoint when the job budget is spent (P2), so a run that already succeeded isn't timed out at the router by checkpoint latency. --- api/src/job.ts | 51 +++++++++++-------- api/src/session-workspace.test.ts | 5 ++ api/src/session-workspace.ts | 18 ++++--- docs/lambda-microvm/README.md | 4 +- service/src/config.ts | 3 ++ .../src/runtime-session/lambda-client-aws.ts | 1 + service/src/runtime-session/lambda-client.ts | 3 ++ service/src/runtime-session/registry.ts | 21 +++++--- service/src/sandbox-backend/index.ts | 1 + .../sandbox-backend/lambda-microvm.test.ts | 16 ++++++ service/src/sandbox-backend/lambda-microvm.ts | 45 ++++++++++++---- 11 files changed, 121 insertions(+), 47 deletions(-) diff --git a/api/src/job.ts b/api/src/job.ts index e40c424..14719dc 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -850,7 +850,12 @@ export class Job { 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) this.session.markPrimed(name, file.id); + 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. */ + this.session.markPrimed(name, file.id, this.inputFileHashes.get(name)?.readOnly === true); + } } private async reusePrimedInput(file: TFile): Promise { @@ -1529,14 +1534,12 @@ export class Job { } let size: number; - let mtimeMs: number; try { /* 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. */ const st = await fsp.lstat(fullPath); size = st.size; - mtimeMs = st.mtimeMs; } catch (err) { this.log.debug({ path: relativePath, err }, 'walkDir: unable to stat file'); return { collected: false, truncated: false, stopLoop: false }; @@ -1545,31 +1548,35 @@ export class Job { return { collected: false, truncated: false, stopLoop: false }; } - /* Session mode output diffing: a persistent workspace still holds prior - * jobs' outputs. Skip any file already surfaced to the client whose - * size+mtime are unchanged, so only genuine deltas are re-uploaded. An - * input file (tracked in inputFileHashes) is left to the existing - * modified/unchanged classification below. */ - const outputSignature = `${size}:${mtimeMs}`; - if (this.session && !this.inputFileHashes.has(relativePath) - && this.session.isSurfaced(relativePath, outputSignature)) { - 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 }; + } + + 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, @@ -1596,7 +1603,9 @@ export class Job { /* 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) this.pendingSurfaced.set(newId, { name: relativePath, signature: outputSignature }); + if (this.session && contentHash != null) { + this.pendingSurfaced.set(newId, { name: relativePath, signature: contentHash }); + } return { collected: true, truncated: false, stopLoop: false }; } diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index 8e4c1a5..c614f30 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -109,6 +109,11 @@ describe('SessionWorkspace state', () => { 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(); + await ws.reset(); expect(ws.isSurfaced('out.csv', '10:100')).toBe(false); expect(ws.primedInputId('in.csv')).toBeUndefined(); diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index 21b2282..d284a10 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -91,9 +91,11 @@ export class SessionWorkspace { * later job re-scanning the persistent workspace does not re-upload * unchanged prior outputs (output diffing). */ private readonly surfaced = new Map(); - /** relPath -> storage file id already primed onto disk, so an unchanged - * input delivered again is not re-downloaded (priming dedup). */ - private readonly primed = new Map(); + /** relPath -> {id, readOnly} 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. */ + private readonly primed = new Map(); constructor(binding: SessionBinding) { this.runtimeSessionId = binding.runtimeSessionId; @@ -132,12 +134,16 @@ export class SessionWorkspace { 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 { - return this.primed.get(relPath); + const entry = this.primed.get(relPath); + if (!entry || entry.readOnly) return undefined; + return entry.id; } - markPrimed(relPath: string, storageFileId: string): void { - this.primed.set(relPath, storageFileId); + markPrimed(relPath: string, storageFileId: string, readOnly = false): void { + this.primed.set(relPath, { id: storageFileId, readOnly }); } /** Full teardown: wipe the dir, release the pinned UID, clear diff state. */ diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md index 4818743..475c357 100644 --- a/docs/lambda-microvm/README.md +++ b/docs/lambda-microvm/README.md @@ -146,6 +146,7 @@ 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 @@ -191,7 +192,8 @@ All names as they appear in `service/src/config.ts`. |---|---|---| | `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 for runtime VM stdout to reach CloudWatch. | +| `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. | diff --git a/service/src/config.ts b/service/src/config.ts index 7814119..f37b02c 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -175,6 +175,9 @@ export const env = { 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), diff --git a/service/src/runtime-session/lambda-client-aws.ts b/service/src/runtime-session/lambda-client-aws.ts index 7798d90..0ee48b6 100644 --- a/service/src/runtime-session/lambda-client-aws.ts +++ b/service/src/runtime-session/lambda-client-aws.ts @@ -109,6 +109,7 @@ export class AwsLambdaMicrovmClient implements LambdaMicrovmClient { autoResumeEnabled: args.idlePolicy.autoResume, } : undefined, + logging: args.logGroup ? { cloudWatch: { logGroup: args.logGroup } } : undefined, runHookPayload: args.runHookPayload, clientToken: args.clientToken, })); diff --git a/service/src/runtime-session/lambda-client.ts b/service/src/runtime-session/lambda-client.ts index e5a966d..3b55cde 100644 --- a/service/src/runtime-session/lambda-client.ts +++ b/service/src/runtime-session/lambda-client.ts @@ -37,6 +37,9 @@ export interface RunMicrovmArgs { 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. */ diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index bd2608a..7f3f23e 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -46,17 +46,22 @@ const LOCK_PREFIX = 'rtsx:lock:'; const GEN_PREFIX = 'rtsx:gen:'; const ACTIVE_ZSET = 'rtsx:active'; -/** The session lock is held across the whole `executeSession` critical path — - * launch + health + execute + post-run checkpoint — so its TTL must outlive the - * sum of those configurable budgets, else a live holder is fenced mid-operation - * and a second worker can acquire the same session concurrently. Derive it from - * the actual budgets (not a fixed placeholder) so raising any one of them keeps - * the lock safe, plus headroom for lock-wait + scheduling jitter. */ +/** The session lock is held across the WHOLE `executeSession` critical path, so + * its TTL must outlive the worst-case sum of the configurable budgets — else a + * live holder is fenced mid-operation and a second worker acquires the same + * session concurrently. A relaunch path can spend, back to back: + * - launch throttle wait (acquireOpBudget) ≤ LAUNCH_TIMEOUT_MS + * - poll to RUNNING (waitUntilRunning) ≤ LAUNCH_TIMEOUT_MS + * - restore the checkpoint before the first exec ≤ CHECKPOINT_TIMEOUT_MS + * - health check ≤ HEALTH_TIMEOUT_MS + * - the execute ≤ JOB_TIMEOUT + * - the post-run checkpoint ≤ CHECKPOINT_TIMEOUT_MS + * so budget 2× launch and 2× checkpoint, plus headroom for lock-wait + jitter. */ export const RUNTIME_SESSION_LOCK_TTL_MS = env.JOB_TIMEOUT + - env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + + 2 * env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS + - env.CHECKPOINT_TIMEOUT_MS + + 2 * env.CHECKPOINT_TIMEOUT_MS + 60_000; const MAX_MICROVM_DURATION_SECONDS = 28_800; diff --git a/service/src/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts index 0c84230..bd47b18 100644 --- a/service/src/sandbox-backend/index.ts +++ b/service/src/sandbox-backend/index.ts @@ -23,6 +23,7 @@ function createBackend(): SandboxBackend { 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, diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 5038b0c..f63423a 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -343,6 +343,22 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { expect(record?.state).toBe('RUNNING'); }); + 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); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index c6caf51..78dd5f6 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -42,6 +42,7 @@ export interface LambdaMicrovmBackendConfig { imageArn: string; imageVersion?: string; executionRoleArn?: string; + logGroup?: string; ingressConnectorArns?: string[]; egressConnectorArns?: string[]; port: number; @@ -161,6 +162,11 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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 }); @@ -181,9 +187,17 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { * 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); + const canCheckpoint = !ctx.signal.aborted && remainingBudgetMs > this.config.checkpoint.timeoutMs; const settled = await readRuntimeSessionRecord(runtimeSessionId); const nextRecord = settled - ? await this.checkpointUnderLock(client, settled, runtimeSessionId, now, lockToken) + ? 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); @@ -219,16 +233,17 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { try { return await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); } catch (error) { - /* Keep the warm VM only when the runner actually responded — a non-2xx - * makes axios throw an AxiosError carrying `.response`, so the VM is alive - * and just the request failed. No response (connection/timeout/abort) or a - * failed health check means the VM is unreachable/dirty: terminate it and - * drop the record so the next call relaunches + restores. (`Error from - * sandbox` covers proxyExecute's manual 2xx-but-not-200 throw.) */ - const sandboxResponded = - (axios.isAxiosError(error) && error.response != null) || - (error instanceof Error && error.message === 'Error from sandbox'); - if (ctx.signal.aborted || !sandboxResponded) { + /* 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`): health/execute + * couldn't reach the VM (connection refused/timeout). + * Everything else keeps the warm VM: a non-2xx sandbox response (AxiosError + * WITH `.response` — the VM is alive, only the request failed) and, crucially, + * a pre-request control-plane failure like a throttled CreateMicrovmAuthToken + * (not an axios error at all) — the VM was never touched. */ + const transportFailure = axios.isAxiosError(error) && error.response == null; + if (ctx.signal.aborted || transportFailure) { await this.terminate(client, vm.microvmId, ctx.signal.aborted ? 'timeout' : 'error').catch(() => {}); await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); } @@ -252,11 +267,18 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { && record.image_arn === this.config.imageArn && record.image_version === this.config.imageVersion && record.port === this.config.port; + /* 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 @@ -431,6 +453,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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, From 06bbb17f5d6a419cc72be5bf70de4a00cf02990f Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 15:00:51 -0400 Subject: [PATCH 17/24] fix: address Codex review round 3 of stateful session mechanics - Tear down the warm session VM on health-check failure, not just on transport errors: assertHealthy wraps failures as MICROVM_UNHEALTHY, and executeOnSessionVm now recycles the VM + drops the registry record for it. - Restore whenever checkpoints are active on relaunch (drop the record-present precondition) so a fresh VM still pulls the last checkpoint. - Bound the checkpoint object-store put with the checkpoint timeout so a stalled S3/MinIO write can't hold the session lock past JOB_TIMEOUT. - Wipe the workspace to a clean slate when a restore fails mid-extract, so the job never runs against a half-applied checkpoint. - Throttle CreateMicrovmAuthToken under a shared per-second token budget (LAMBDA_MICROVM_TOKEN_TPS, default 8) so concurrent warm-session executes queue instead of bursting past the AWS TPS limit, mirroring launch's run budget. --- api/src/session-checkpoint.ts | 6 ++ service/src/config.ts | 4 ++ service/src/runtime-session/checkpoint.ts | 22 +++++++- service/src/sandbox-backend/index.ts | 1 + .../sandbox-backend/lambda-microvm.test.ts | 1 + service/src/sandbox-backend/lambda-microvm.ts | 56 ++++++++++++++----- 6 files changed, 74 insertions(+), 16 deletions(-) diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index 8137044..31c9fbc 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -80,6 +80,12 @@ export async function restoreSessionCheckpoint(req: Request, res: Response): Pro 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' }); } } diff --git a/service/src/config.ts b/service/src/config.ts index f37b02c..5f0c30e 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -194,6 +194,10 @@ export const env = { 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 diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts index e11e9a9..81f36a8 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -11,6 +11,19 @@ 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 @@ -126,7 +139,14 @@ export async function checkpointSession(args: { microvmCheckpoints.inc({ outcome: 'skipped_busy' }); return 'skipped_busy'; } - await args.store.put(args.runtimeSessionId, data); + /* Bound the object-store write by the checkpoint timeout too — otherwise a + * stalled S3/MinIO put holds the session lock past JOB_TIMEOUT (and, if it + * outlives the lock TTL, could clobber a newer checkpoint). */ + await withTimeout( + args.store.put(args.runtimeSessionId, data), + args.config.timeoutMs, + 'checkpoint store.put', + ); microvmCheckpointBytes.observe(data.length); microvmCheckpoints.inc({ outcome: 'stored' }); return 'stored'; diff --git a/service/src/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts index bd47b18..481ee4c 100644 --- a/service/src/sandbox-backend/index.ts +++ b/service/src/sandbox-backend/index.ts @@ -32,6 +32,7 @@ function createBackend(): SandboxBackend { 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, diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index f63423a..aca5367 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -111,6 +111,7 @@ function config(overrides: Partial = {}): LambdaMicr launchTimeoutMs: 2_000, healthTimeoutMs: 1_000, launchTps: 50, + tokenTps: 50, jobTimeoutMs: 300_000, idleSeconds: 300, suspendedSeconds: 1_800, diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 78dd5f6..121435f 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -1,6 +1,6 @@ import axios from 'axios'; import { nanoid } from 'nanoid'; -import type { LambdaMicrovmClient, MicrovmDescription, MicrovmIdlePolicy } from '../runtime-session/lambda-client'; +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'; @@ -51,6 +51,7 @@ export interface LambdaMicrovmBackendConfig { launchTimeoutMs: number; healthTimeoutMs: number; launchTps: number; + tokenTps: number; jobTimeoutMs: number; /* Session-mode (find-or-launch) tuning. */ idleSeconds: number; @@ -236,14 +237,17 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { /* 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`): health/execute - * couldn't reach the VM (connection refused/timeout). + * - 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, crucially, - * a pre-request control-plane failure like a throttled CreateMicrovmAuthToken + * 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. */ const transportFailure = axios.isAxiosError(error) && error.response == null; - if (ctx.signal.aborted || transportFailure) { + const unhealthy = error instanceof SandboxBackendError && error.code === 'MICROVM_UNHEALTHY'; + if (ctx.signal.aborted || transportFailure || unhealthy) { await this.terminate(client, vm.microvmId, ctx.signal.aborted ? 'timeout' : 'error').catch(() => {}); await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); } @@ -330,10 +334,13 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { throw new SandboxBackendError('MICROVM_FENCED', `Lost session lock for ${runtimeSessionId} after launch`); } - /* Fresh VM for an existing session: restore its predecessor's workspace - * before the first execute, so an 8h rollover / eviction is invisible. - * A prior record (or a checkpoint pointer) means the session existed. */ - if (this.checkpointStore && this.checkpointsActive() && (record?.workspace_checkpoint || record != null)) { + /* 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()) { await restoreSession({ client, store: this.checkpointStore, @@ -379,6 +386,29 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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). */ + 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; + } + return client.createMicrovmAuthToken({ + microvmId, + port: this.config.port, + ttlSeconds: this.config.authTokenTtlSeconds, + }); + } + private async proxyExecute( client: LambdaMicrovmClient, vm: MicrovmDescription, @@ -387,11 +417,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { runtimeSessionId?: string, ): Promise { const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); - const token = await client.createMicrovmAuthToken({ - microvmId: vm.microvmId, - port: this.config.port, - ttlSeconds: this.config.authTokenTtlSeconds, - }); + const token = await this.mintAuthToken(client, vm.microvmId); await this.assertHealthy(base, token.token, ctx); /* Session mode is opted into per-request via this header (not a /run From 78673a2fa86e234961c5dbfb366f928d14101231 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 15:48:29 -0400 Subject: [PATCH 18/24] fix: address Codex review round 4 + Macroscope of stateful sessions Codex round 4 (7) and the open Macroscope batch (3): - Preserve priming/output-diff state across a checkpoint restore (P1). A relaunched VM started with an empty primed map and re-downloaded every input ref, overwriting a restored in-place-modified file with its original. The checkpoint now carries a sidecar (SessionWorkspace.snapshotMeta) that the restore rebuilds (loadMeta), so warm-VM behavior survives a relaunch. - Version checkpoint objects by a per-checkpoint monotonic sequence and read the highest on restore. A put that timed out and lands late writes an older key and can never overwrite a newer checkpoint; restore still works with no Redis record (list by session prefix). Best-effort prune of strictly-older keys. - mintAuthToken maps control-plane failures instead of letting them escape raw: throttled -> poisonOpBucket('token') + MICROVM_LAUNCH_THROTTLED; not_found (VM evicted) -> MICROVM_UNHEALTHY so the caller tears down the stale record and relaunches; other -> MICROVM_LAUNCH_FAILED. - Route checkpoint/restore token mints through the same token budget via a mintToken callback, so a burst of concurrent sessions can't bypass the TPS cap. - Bound the checkpoint fetch (store.get) with the checkpoint timeout, symmetric with the put, so a stalled S3 read can't hold the session lock through a relaunch. - Terminate a superseded (config/version/port drift, deadline-near) session VM before relaunching, except when AWS already auto-terminated it, so it isn't left running and billing. - Restore extracts with --strip-components=1 -C dir so a poisoned archive member can't escape the workspace into shared runner space (Macroscope corrected: a plain -C dir double-nests the session/ prefix). - Fail fast at startup when session checkpoints are enabled without object storage configured, instead of silently falling back to a localhost/test bucket and dropping workspace state on the first relaunch. --- api/src/session-checkpoint.ts | 43 +++++++-- api/src/session-workspace.test.ts | 18 ++++ api/src/session-workspace.ts | 29 +++++++ .../runtime-session/checkpoint-store.test.ts | 45 +++++++--- .../src/runtime-session/checkpoint-store.ts | 87 +++++++++++++++---- service/src/runtime-session/checkpoint.ts | 62 ++++++------- service/src/runtime-session/registry.ts | 12 +++ .../sandbox-backend/lambda-microvm.test.ts | 34 +++++++- service/src/sandbox-backend/lambda-microvm.ts | 45 ++++++++-- service/src/secure-startup.test.ts | 15 ++++ service/src/secure-startup.ts | 19 ++++ 11 files changed, 330 insertions(+), 79 deletions(-) diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index 31c9fbc..5629d01 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -5,7 +5,8 @@ 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 { getBoundSessionWorkspace } from './session-workspace'; +import type { SessionMetaSnapshot, SessionWorkspace } from './session-workspace'; +import { SESSION_META_FILE, getBoundSessionWorkspace } from './session-workspace'; /** * Session workspace checkpoint / restore. @@ -31,7 +32,14 @@ export async function streamSessionCheckpoint(res: Response): Promise { res.status(409).json({ message: 'No session workspace is bound' }); return; } - await session.ownership(); + 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 from the live + * workspace once tar has read it. */ + const metaPath = path.join(dir, SESSION_META_FILE); + await fsp.writeFile(metaPath, JSON.stringify(session.snapshotMeta())); res.status(200); res.setHeader('Content-Type', CHECKPOINT_CONTENT_TYPE); @@ -47,6 +55,8 @@ export async function streamSessionCheckpoint(res: Response): Promise { logger.error({ err: error }, 'Failed to stream session checkpoint'); if (!res.headersSent) res.status(500).json({ message: 'checkpoint failed' }); else res.destroy(); + } finally { + await fsp.rm(metaPath, { force: true }).catch(() => {}); } } @@ -65,10 +75,12 @@ export async function restoreSessionCheckpoint(req: Request, res: Response): Pro await fsp.mkdir(dir, { recursive: true }); /* Archives are created with relative `session/...` members (see the create - * side), so extraction stays within the workspace root without needing - * GNU-only flags. Production hardening: verify/scan the archive before - * trusting a restore from shared storage. */ - const tar = spawn('tar', ['-xzf', '-', '-C', SANDBOX_WORKSPACE_ROOT], { + * 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')); @@ -77,6 +89,7 @@ export async function restoreSessionCheckpoint(req: Request, res: Response): Pro 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'); @@ -90,6 +103,24 @@ export async function restoreSessionCheckpoint(req: Request, res: Response): Pro } } +/** 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 { + const parsed = JSON.parse(await fsp.readFile(metaPath, 'utf8')) as SessionMetaSnapshot; + if (Array.isArray(parsed?.primed) && Array.isArray(parsed?.surfaced)) { + session.loadMeta(parsed); + } + } catch (error) { + logger.debug({ err: error }, 'No session meta sidecar to restore'); + } finally { + await fsp.rm(metaPath, { force: true }).catch(() => {}); + } +} + 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 }); diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index c614f30..c7190ab 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -122,4 +122,22 @@ describe('SessionWorkspace state', () => { 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'); + 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); + /* 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 index d284a10..02011e8 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -36,6 +36,17 @@ import { /** 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'; + +export interface SessionMetaSnapshot { + primed: Array<[string, { id: string; readOnly: boolean }]>; + surfaced: Array<[string, string]>; +} + const RUNTIME_SESSION_ID_PATTERN = /^[A-Za-z0-9._:-]{1,128}$/; export interface SessionBinding { @@ -146,6 +157,24 @@ export class SessionWorkspace { this.primed.set(relPath, { id: storageFileId, readOnly }); } + /** 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(); diff --git a/service/src/runtime-session/checkpoint-store.test.ts b/service/src/runtime-session/checkpoint-store.test.ts index 635945e..1fd5064 100644 --- a/service/src/runtime-session/checkpoint-store.test.ts +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -1,19 +1,27 @@ import { describe, expect, test } from 'bun:test'; -import { MemoryCheckpointStore, CheckpointTooLargeError, checkpointObjectKey } from './checkpoint-store'; +import { + MemoryCheckpointStore, + CheckpointTooLargeError, + checkpointObjectKey, + checkpointPrefixFor, +} from './checkpoint-store'; const BIG = 1_000_000; describe('checkpoint store', () => { - test('object key is deterministic per runtime session under the prefix', () => { - expect(checkpointObjectKey('rt_abc')).toBe('rtsx-checkpoints/rt_abc.tar.gz'); - expect(checkpointObjectKey('rt_abc')).toBe(checkpointObjectKey('rt_abc')); - expect(checkpointObjectKey('rt_xyz')).not.toBe(checkpointObjectKey('rt_abc')); + 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 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 bytes and copies defensively', async () => { + test('memory store round-trips the latest bytes and copies defensively', async () => { const store = new MemoryCheckpointStore(); const original = Buffer.from('workspace-bytes'); - await store.put('rt_1', original); + await store.put('rt_1', 1, original); const fetched = await store.get('rt_1', BIG); expect(fetched?.toString()).toBe('workspace-bytes'); @@ -27,16 +35,29 @@ describe('checkpoint store', () => { expect(await store.get('rt_missing', BIG)).toBeNull(); }); - test('last-writer-wins on the same key', async () => { + test('get reads the highest sequence, not the last write', async () => { const store = new MemoryCheckpointStore(); - await store.put('rt_1', Buffer.from('v1')); - await store.put('rt_1', Buffer.from('v2')); - expect((await store.get('rt_1', BIG))?.toString()).toBe('v2'); + await store.put('rt_1', 2, Buffer.from('newer')); + /* a stale 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('put prunes 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')); + /* only the newest object survives; siblings for other sessions are untouched */ + 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', Buffer.alloc(2048)); + 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 index f09fd47..76219a2 100644 --- a/service/src/runtime-session/checkpoint-store.ts +++ b/service/src/runtime-session/checkpoint-store.ts @@ -2,24 +2,34 @@ import { Client as MinioClient } from 'minio'; import { env } from '../config'; /** - * Durable storage for session workspace checkpoints. Deterministic key per - * runtime session (`.tar.gz`) so a relaunch can - * find the latest checkpoint even if the registry record was lost. Writes are - * serialized by the session lock, so last-writer-wins is the intended - * semantic; object versioning (Phase 4) adds forensic history on top. + * Durable storage for session workspace checkpoints. Each checkpoint writes a + * distinct, strictly increasing object under a per-session prefix + * (`/.tar.gz`); restore reads the highest + * sequence. A put that timed out and lands late therefore writes an OLDER + * sequence and can never overwrite a newer checkpoint, and a relaunch can still + * find the latest checkpoint by listing the prefix even if the registry record + * was lost. Older sequences are best-effort pruned after each successful put. */ export interface CheckpointStore { - put(runtimeSessionId: string, data: Buffer): Promise; - /** `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. */ + 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; } export class CheckpointTooLargeError extends Error {} -export function checkpointObjectKey(runtimeSessionId: string): string { - return `${env.CHECKPOINT_PREFIX}${runtimeSessionId}.tar.gz`; +/** 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`; } /** S3/MinIO-backed store using the same MINIO_* envs as file-server. */ @@ -40,14 +50,20 @@ export class MinioCheckpointStore implements CheckpointStore { this.bucket = process.env.CODEAPI_CHECKPOINT_BUCKET ?? process.env.MINIO_BUCKET ?? 'test-bucket'; } - async put(runtimeSessionId: string, data: Buffer): Promise { - await this.client.putObject(this.bucket, checkpointObjectKey(runtimeSessionId), data, data.length, { + 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 stale late put must not delete the checkpoint that superseded it — + * restore always reads the max sequence). */ + await this.pruneOlderThan(runtimeSessionId, key).catch(() => {}); } async get(runtimeSessionId: string, maxBytes: number): Promise { - const key = checkpointObjectKey(runtimeSessionId); + const key = await this.latestKey(runtimeSessionId); + if (!key) return null; let size: number; try { const stat = await this.client.statObject(this.bucket, key); @@ -75,19 +91,52 @@ export class MinioCheckpointStore implements CheckpointStore { } 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. */ +/** 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, data: Buffer): Promise { - this.objects.set(runtimeSessionId, Buffer.from(data)); + 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 sequences only — never a newer one */ + if (existing.startsWith(prefix) && existing < key) this.objects.delete(existing); + } + this.objects.set(key, Buffer.from(data)); } async get(runtimeSessionId: string, maxBytes: number): Promise { - const data = this.objects.get(runtimeSessionId); - if (!data) return null; + 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`); } diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts index 81f36a8..3d970bf 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -1,8 +1,9 @@ import axios from 'axios'; -import type { LambdaMicrovmClient } from './lambda-client'; +import type { MicrovmAuthToken } from './lambda-client'; import type { CheckpointStore } from './checkpoint-store'; import { acquireRuntimeSessionLock, + allocateCheckpointSequence, readRuntimeSessionRecord, releaseRuntimeSessionLock, writeRuntimeSessionRecord, @@ -50,15 +51,10 @@ export interface CheckpointConfig { const RUNTIME_SESSION_ID_HEADER = 'X-Runtime-Session-Id'; export async function pullCheckpoint( - client: LambdaMicrovmClient, - args: { microvmId: string; endpointBase: string; runtimeSessionId: string }, + args: { mintToken: () => Promise; endpointBase: string; runtimeSessionId: string }, config: CheckpointConfig, ): Promise { - const token = await client.createMicrovmAuthToken({ - microvmId: args.microvmId, - port: config.port, - ttlSeconds: config.authTokenTtlSeconds, - }); + const token = await args.mintToken(); const response = await axios.get(`${args.endpointBase}/api/v2/session/checkpoint`, { headers: { [token.headerName]: token.token, @@ -72,16 +68,11 @@ export async function pullCheckpoint( } export async function pushRestore( - client: LambdaMicrovmClient, - args: { microvmId: string; endpointBase: string; runtimeSessionId: string }, + args: { mintToken: () => Promise; endpointBase: string; runtimeSessionId: string }, data: Buffer, config: CheckpointConfig, ): Promise { - const token = await client.createMicrovmAuthToken({ - microvmId: args.microvmId, - port: config.port, - ttlSeconds: config.authTokenTtlSeconds, - }); + const token = await args.mintToken(); await axios.post(`${args.endpointBase}/api/v2/session/restore`, data, { headers: { [token.headerName]: token.token, @@ -102,7 +93,7 @@ export async function pushRestore( * post-checkpoint will cover this one. */ export async function checkpointSession(args: { - client: LambdaMicrovmClient; + mintToken: (microvmId: string) => Promise; store: CheckpointStore; runtimeSessionId: string; config: CheckpointConfig; @@ -121,18 +112,20 @@ export async function checkpointSession(args: { microvmCheckpoints.inc({ outcome: 'skipped_state' }); return 'skipped_state'; } - const data = await pullCheckpoint(args.client, { - microvmId: record.microvm_id, + const microvmId = record.microvm_id; + const data = await pullCheckpoint({ + mintToken: () => args.mintToken(microvmId), endpointBase: args.normalizeEndpoint(record.endpoint), runtimeSessionId: args.runtimeSessionId, }, args.config); - /* Fence BEFORE writing the object store. The checkpoint key is - * deterministic per session (last-writer-wins), so a caller whose lock - * expired must not clobber a newer blob. The fenced record write is the - * ownership check: if it reports we were fenced, skip the store entirely. */ + /* Each checkpoint writes a distinct, strictly increasing object key, so a + * put that stalled past the lock and lands late writes an OLDER key and can + * never overwrite the newer one restore reads. Fence the pointer write + * first: if it reports we were fenced, skip the store entirely. */ + const sequence = await allocateCheckpointSequence(args.runtimeSessionId); const persisted = await writeRuntimeSessionRecord({ ...record, - workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId), + workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId, sequence), checkpointed_at: Date.now(), }, lockToken); if (!persisted) { @@ -140,10 +133,9 @@ export async function checkpointSession(args: { 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 (and, if it - * outlives the lock TTL, could clobber a newer checkpoint). */ + * stalled S3/MinIO put holds the session lock past JOB_TIMEOUT. */ await withTimeout( - args.store.put(args.runtimeSessionId, data), + args.store.put(args.runtimeSessionId, sequence, data), args.config.timeoutMs, 'checkpoint store.put', ); @@ -165,7 +157,7 @@ export async function checkpointSession(args: { /** Relaunch restore: caller holds the session lock and the VM is RUNNING. */ export async function restoreSession(args: { - client: LambdaMicrovmClient; + mintToken: (microvmId: string) => Promise; store: CheckpointStore; runtimeSessionId: string; microvmId: string; @@ -173,10 +165,16 @@ export async function restoreSession(args: { config: CheckpointConfig; }): Promise<'restored' | 'absent' | '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. */ + * 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 args.store.get(args.runtimeSessionId, args.config.maxBytes); + data = await withTimeout( + args.store.get(args.runtimeSessionId, args.config.maxBytes), + args.config.timeoutMs, + 'checkpoint store.get', + ); } catch (error) { microvmRestores.inc({ outcome: 'failed' }); logger.warn('Checkpoint fetch failed; continuing with a fresh workspace', { @@ -190,7 +188,11 @@ export async function restoreSession(args: { return 'absent'; } try { - await pushRestore(args.client, { microvmId: args.microvmId, endpointBase: args.endpointBase, runtimeSessionId: args.runtimeSessionId }, data, args.config); + 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, diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index 7f3f23e..d57effa 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -44,6 +44,7 @@ export interface RuntimeSessionRecord { 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, so @@ -201,6 +202,17 @@ export async function allocateRuntimeSessionGeneration(runtimeSessionId: string) return generation; } +/** Per-checkpoint monotonic sequence. Each successful checkpoint writes a + * distinct, strictly increasing object key so a put that timed out and lands + * late can never overwrite a newer checkpoint (restore always reads the max + * sequence). Bumps on every checkpoint, unlike the per-relaunch generation. */ +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; +} + export async function touchRuntimeSessionActive(runtimeSessionId: string, lastSeenAtMs: number): Promise { await redis.zadd(ACTIVE_ZSET, lastSeenAtMs, runtimeSessionId); } diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index aca5367..998923b 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -412,6 +412,32 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { expect(fake.callsFor('terminateMicrovm')).toHaveLength(1); expect(await readRuntimeSessionRecord('rt_session_1')).toBeNull(); }); + + 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); + }); }); describe('LambdaMicrovmSandboxBackend auto-checkpoint', () => { @@ -438,20 +464,20 @@ describe('LambdaMicrovmSandboxBackend auto-checkpoint', () => { /* 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(store.objects.get('rt_ckpt_1')?.toString()).toBe(checkpointBlob); + expect((await store.get('rt_ckpt_1', 1_000_000))?.toString()).toBe(checkpointBlob); const record = await readRuntimeSessionRecord('rt_ckpt_1'); - expect(record?.workspace_checkpoint).toBe(checkpointObjectKey('rt_ckpt_1')); + expect(record?.workspace_checkpoint).toBe(checkpointObjectKey('rt_ckpt_1', 1)); 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', Buffer.from('PRIOR_WORKSPACE')); + await store.put('rt_ckpt_1', 1, 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'), + state: 'TERMINATED', generation: 3, last_seen_at: 1, workspace_checkpoint: checkpointObjectKey('rt_ckpt_1', 1), }, seedToken as string); const { releaseRuntimeSessionLock } = await import('../runtime-session/registry'); await releaseRuntimeSessionLock('rt_ckpt_1', seedToken as string); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 121435f..0941104 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -290,6 +290,16 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { return { microvmId: record.microvm_id as string, state: 'RUNNING', endpoint: record.endpoint }; } + /* 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, @@ -342,7 +352,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { * just no-ops one stat). */ if (this.checkpointStore && this.checkpointsActive()) { await restoreSession({ - client, + mintToken: (microvmId) => this.mintAuthToken(client, microvmId), store: this.checkpointStore, runtimeSessionId, microvmId: vm.microvmId, @@ -371,7 +381,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { return base; } const result = await checkpointSession({ - client, + mintToken: (microvmId) => this.mintAuthToken(client, microvmId), store: this.checkpointStore, runtimeSessionId, config: this.config.checkpoint, @@ -388,7 +398,11 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { /** 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). */ + * 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', { @@ -402,11 +416,26 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { } throw error; } - return client.createMicrovmAuthToken({ - microvmId, - port: this.config.port, - ttlSeconds: this.config.authTokenTtlSeconds, - }); + 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( diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 39d3bcf..705c8ae 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -146,6 +146,11 @@ describe('sandbox backend policy', () => { 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', () => { @@ -185,6 +190,16 @@ describe('sandbox backend policy', () => { 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; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 0bcd1b2..8093d77 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -91,6 +91,25 @@ export function validateSandboxBackendPolicy(): void { '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 { From 2f75eff7fa7ea4531babffa85a373983dd4389e3 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 09:32:02 -0400 Subject: [PATCH 19/24] fix: address Codex review round 5 of stateful session mechanics - Refuse a symlinked checkpoint sidecar (P1). The round-4 primed-persistence sidecar was written with a plain (root) writeFile; sandboxed code could squat .codeapi-session-meta.json as a symlink and turn the post-exec checkpoint into an arbitrary-file overwrite. Unlink first (never follows a link) then create a fresh regular file exclusively (wx); on restore, lstat-guard the read so a symlinked sidecar in an untrusted archive is never followed. - Route non-8080 MicroVM ports with X-aws-proxy-port. Health/execute/checkpoint traffic omitted the port header, so a non-default LAMBDA_MICROVM_PORT would be routed to 8080 (the token is minted for the configured port). Added microvmPortHeaders(port) and applied it to all proxied requests. - Reserve the whole checkpoint path budget before starting it. The skip guard checked only one checkpoint timeout, but checkpointUnderLock can spend a token-budget wait + GET + put; a run finishing with barely more than one timeout left could still blow waitUntilFinished(JOB_TIMEOUT). Require launchTimeoutMs + 2*checkpoint.timeoutMs of headroom. - Baseline reused writable inputs against the original upload hash, not a re-hash of the on-disk copy. A prior turn's in-place mutation was banked as pristine, letting the walker echo the original ref as unchanged. The primed map now retains the original hash (round-trips through the checkpoint sidecar) and reuse baselines against it, so a mutation is reported modified-from- original without re-downloading (preserves the round-4 persistence fix). --- api/src/job.ts | 13 +++++++--- api/src/session-checkpoint.ts | 14 +++++++++-- api/src/session-workspace.test.ts | 5 +++- api/src/session-workspace.ts | 25 +++++++++++++------ service/src/runtime-session/checkpoint.ts | 3 +++ service/src/runtime-session/lambda-client.ts | 10 ++++++++ .../sandbox-backend/lambda-microvm.test.ts | 13 ++++++++++ service/src/sandbox-backend/lambda-microvm.ts | 15 ++++++++--- 8 files changed, 81 insertions(+), 17 deletions(-) diff --git a/api/src/job.ts b/api/src/job.ts index 14719dc..5368bb7 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -853,8 +853,10 @@ export class Job { 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. */ - this.session.markPrimed(name, file.id, this.inputFileHashes.get(name)?.readOnly === true); + * 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); } } @@ -866,7 +868,12 @@ export class Job { try { const st = await fsp.lstat(filePath); if (!st.isFile()) return false; - const hash = await this.computeFileHash(filePath, true); + /* 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, diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index 5629d01..7570157 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -37,9 +37,14 @@ export async function streamSessionCheckpoint(res: Response): Promise { /* 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 from the live - * workspace once tar has read it. */ + * workspace once tar has read it. Sandboxed code from a prior exec can squat + * this name as a symlink; the API process runs as root, so unlink it first + * (unlink never follows a link) then create a fresh regular file exclusively + * (`wx`) — otherwise a privileged write would follow the link and clobber an + * arbitrary target outside the workspace. */ const metaPath = path.join(dir, SESSION_META_FILE); - await fsp.writeFile(metaPath, JSON.stringify(session.snapshotMeta())); + await fsp.rm(metaPath, { force: true }); + await fsp.writeFile(metaPath, JSON.stringify(session.snapshotMeta()), { flag: 'wx' }); res.status(200); res.setHeader('Content-Type', CHECKPOINT_CONTENT_TYPE); @@ -110,6 +115,11 @@ export async function restoreSessionCheckpoint(req: Request, res: Response): Pro 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; if (Array.isArray(parsed?.primed) && Array.isArray(parsed?.surfaced)) { session.loadMeta(parsed); diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index c7190ab..c79367e 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -126,7 +126,7 @@ describe('SessionWorkspace state', () => { 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'); + 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 @@ -137,6 +137,9 @@ describe('SessionWorkspace state', () => { 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 index 02011e8..de72d5c 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -43,7 +43,7 @@ export const RUNTIME_SESSION_ID_HEADER = 'x-runtime-session-id'; export const SESSION_META_FILE = '.codeapi-session-meta.json'; export interface SessionMetaSnapshot { - primed: Array<[string, { id: string; readOnly: boolean }]>; + primed: Array<[string, { id: string; readOnly: boolean; hash?: string }]>; surfaced: Array<[string, string]>; } @@ -102,11 +102,14 @@ export class SessionWorkspace { * later job re-scanning the persistent workspace does not re-upload * unchanged prior outputs (output diffing). */ private readonly surfaced = new Map(); - /** relPath -> {id, readOnly} 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. */ - private readonly primed = 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; @@ -153,8 +156,14 @@ export class SessionWorkspace { return entry.id; } - markPrimed(relPath: string, storageFileId: string, readOnly = false): void { - this.primed.set(relPath, { id: storageFileId, readOnly }); + 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 diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts index 3d970bf..c1c896c 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -1,5 +1,6 @@ import axios from 'axios'; import type { MicrovmAuthToken } from './lambda-client'; +import { microvmPortHeaders } from './lambda-client'; import type { CheckpointStore } from './checkpoint-store'; import { acquireRuntimeSessionLock, @@ -58,6 +59,7 @@ export async function pullCheckpoint( 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', @@ -76,6 +78,7 @@ export async function pushRestore( 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', }, diff --git a/service/src/runtime-session/lambda-client.ts b/service/src/runtime-session/lambda-client.ts index 3b55cde..676c7e1 100644 --- a/service/src/runtime-session/lambda-client.ts +++ b/service/src/runtime-session/lambda-client.ts @@ -75,6 +75,16 @@ export class LambdaMicrovmApiError extends Error { 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; diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 998923b..32162b9 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -413,6 +413,19 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { 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); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 0941104..99c8cf3 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -5,7 +5,7 @@ import type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, Sandbox import type { RuntimeSessionRecord } from '../runtime-session/registry'; import type { CheckpointConfig } from '../runtime-session/checkpoint'; import type { CheckpointStore } from '../runtime-session/checkpoint-store'; -import { LambdaMicrovmApiError } from '../runtime-session/lambda-client'; +import { LambdaMicrovmApiError, microvmPortHeaders } from '../runtime-session/lambda-client'; import { MicrovmOpThrottledError, acquireOpBudget, poisonOpBucket } from '../runtime-session/throttle'; import { checkpointSession, restoreSession } from '../runtime-session/checkpoint'; import { @@ -193,7 +193,15 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { * 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); - const canCheckpoint = !ctx.signal.aborted && remainingBudgetMs > this.config.checkpoint.timeoutMs; + /* 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 @@ -470,6 +478,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { headers: { ...injectTraceHeaders(req.headers), [token.headerName]: token.token, + ...microvmPortHeaders(this.config.port), ...sessionHeader, }, signal: ctx.signal, @@ -574,7 +583,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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 }, + headers: { 'X-aws-proxy-auth': token, ...microvmPortHeaders(this.config.port) }, timeout: this.config.healthTimeoutMs, signal: ctx.signal, }); From 299a77b47412b18960f19e907f491799e8e5c74f Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 10:10:07 -0400 Subject: [PATCH 20/24] fix: address Codex review round 6 of stateful session mechanics - Drop worker-only backend validation from API-only startup. An API pod authenticates + enqueues and never builds the Lambda backend or checkpoint store, so validateSandboxBackendPolicy() there forced LAMBDA_MICROVM_* and the MINIO_* checkpoint creds into API pods just to boot. Worker/combined startup still validate it. - Map MicroVM poll (GetMicrovm) failures like runMicrovm errors. A throttle or transient control-plane error during waitUntilRunning rethrew a raw LambdaMicrovmApiError -> generic 500; now throttled -> MICROVM_LAUNCH_THROTTLED (+poison), other -> MICROVM_LAUNCH_FAILED. - Don't re-surface previously-primed session inputs as generated outputs. A file primed in an earlier turn persists in the workspace; if a later turn omits it, it has no inputFileHashes/surfaced entry and was echoed as a brand new output. Skip paths the session primed as inputs (SessionWorkspace .isPrimedInput). - Remove a squatted checkpoint metadata directory before writing. The round-5 no-follow rm was non-recursive, so a sandboxed `.codeapi-session-meta.json` directory made every checkpoint fail EISDIR; rm now recursive (still no-follow for symlinks). - Key checkpoint objects by the wall-clock time they were taken, not a Redis sequence counter. The counter carried the record TTL, so after a long idle it reset to 1 while retained S3 objects kept a higher sequence, and restore (which reads the lexicographically-greatest key) restored stale state. A never- resetting timestamp sorts correctly across idle gaps and lets prune drop the stale object; a stalled late put still loses because it carries its earlier start time. --- api/src/job.ts | 8 ++++ api/src/session-checkpoint.ts | 12 ++--- api/src/session-workspace.test.ts | 6 +++ api/src/session-workspace.ts | 7 +++ service/src/lifecycle.ts | 6 ++- .../runtime-session/checkpoint-store.test.ts | 36 ++++++++------- .../src/runtime-session/checkpoint-store.ts | 44 ++++++++++--------- service/src/runtime-session/checkpoint.ts | 20 +++++---- service/src/runtime-session/registry.ts | 12 ----- .../sandbox-backend/lambda-microvm.test.ts | 10 +++-- service/src/sandbox-backend/lambda-microvm.ts | 16 ++++++- 11 files changed, 108 insertions(+), 69 deletions(-) diff --git a/api/src/job.ts b/api/src/job.ts index 5368bb7..c4243c4 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -1578,6 +1578,14 @@ export class Job { 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. It persists in the workspace as an input, not an output. */ + if (this.session && inputFileInfo == null && this.session.isPrimedInput(relativePath)) { + return { collected: false, truncated: false, stopLoop: false }; + } + let wasModified = false; if (inputFileInfo && contentHash != null) { wasModified = contentHash !== inputFileInfo.hash; diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index 7570157..7f67e45 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -38,12 +38,14 @@ export async function streamSessionCheckpoint(res: Response): Promise { * rebuilds it (see restoreSessionCheckpoint). Written under the held session * lock, so no concurrent user code sees it, and removed from the live * workspace once tar has read it. Sandboxed code from a prior exec can squat - * this name as a symlink; the API process runs as root, so unlink it first - * (unlink never follows a link) then create a fresh regular file exclusively - * (`wx`) — otherwise a privileged write would follow the link and clobber an - * arbitrary target outside the workspace. */ + * this name as a symlink OR a directory; the API process runs as root, so + * remove whatever squats the path first — `recursive` clears a directory + * (else the write later fails EISDIR and breaks every checkpoint), `force` + * ignores absence, and neither follows a symlink — then create a fresh + * regular file exclusively (`wx`), so a privileged write can't follow a link + * and clobber an arbitrary target outside the workspace. */ const metaPath = path.join(dir, SESSION_META_FILE); - await fsp.rm(metaPath, { force: true }); + await fsp.rm(metaPath, { force: true, recursive: true }); await fsp.writeFile(metaPath, JSON.stringify(session.snapshotMeta()), { flag: 'wx' }); res.status(200); diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index c79367e..16fd2a1 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -106,6 +106,7 @@ describe('SessionWorkspace state', () => { 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'); @@ -113,6 +114,11 @@ describe('SessionWorkspace state', () => { * (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); await ws.reset(); expect(ws.isSurfaced('out.csv', '10:100')).toBe(false); diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index de72d5c..41ce79d 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -156,6 +156,13 @@ export class SessionWorkspace { 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); + } + markPrimed(relPath: string, storageFileId: string, readOnly = false, hash?: string): void { this.primed.set(relPath, { id: storageFileId, readOnly, hash }); } diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 9b21d05..59a18a8 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -75,7 +75,11 @@ function setupQueueListeners(queue: Queue, name: string): void { export async function startupApiOnly(): Promise { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); - validateSandboxBackendPolicy(); + /* 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) diff --git a/service/src/runtime-session/checkpoint-store.test.ts b/service/src/runtime-session/checkpoint-store.test.ts index 1fd5064..e321df7 100644 --- a/service/src/runtime-session/checkpoint-store.test.ts +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -9,19 +9,20 @@ import { 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'); + test('object key is per session + timestamp, zero-padded for lexical order', () => { + expect(checkpointObjectKey('rt_abc', 1)).toBe('rtsx-checkpoints/rt_abc/000000000000001.tar.gz'); expect(checkpointPrefixFor('rt_abc')).toBe('rtsx-checkpoints/rt_abc/'); - /* zero-padding keeps lexical order == numeric order across widths */ + /* zero-padding keeps lexical order == chronological 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_abc', 1_783_000_000_000) > checkpointObjectKey('rt_abc', 999)).toBe(true); expect(checkpointObjectKey('rt_xyz', 1)).not.toBe(checkpointObjectKey('rt_abc', 1)); }); - test('memory store round-trips the latest bytes and copies defensively', async () => { + test('memory store round-trips the newest bytes and copies defensively', async () => { const store = new MemoryCheckpointStore(); const original = Buffer.from('workspace-bytes'); - await store.put('rt_1', 1, original); + await store.put('rt_1', 1000, original); const fetched = await store.get('rt_1', BIG); expect(fetched?.toString()).toBe('workspace-bytes'); @@ -35,29 +36,30 @@ describe('checkpoint store', () => { expect(await store.get('rt_missing', BIG)).toBeNull(); }); - test('get reads the highest sequence, not the last write', async () => { + test('get reads the newest timestamp, not the last write', async () => { const store = new MemoryCheckpointStore(); - await store.put('rt_1', 2, Buffer.from('newer')); - /* a stale put from a lower sequence lands late but must NOT win */ - await store.put('rt_1', 1, Buffer.from('stale')); + await store.put('rt_1', 2000, Buffer.from('newer')); + /* a stale put carrying an EARLIER start time lands late but must NOT win — + * and a fresh checkpoint after an idle gap (later timestamp) always does */ + await store.put('rt_1', 1000, Buffer.from('stale')); expect((await store.get('rt_1', BIG))?.toString()).toBe('newer'); }); - test('put prunes older sequences for the session', async () => { + test('put prunes older timestamps 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_1', 1000, Buffer.from('v1')); + await store.put('rt_1', 2000, Buffer.from('v2')); /* only the newest object survives; siblings for other sessions are untouched */ - await store.put('rt_other', 1, Buffer.from('other')); + await store.put('rt_other', 1000, 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)); + expect(keys).toContain(checkpointObjectKey('rt_1', 2000)); + expect(keys).not.toContain(checkpointObjectKey('rt_1', 1000)); + expect(keys).toContain(checkpointObjectKey('rt_other', 1000)); }); test('rejects a checkpoint larger than maxBytes', async () => { const store = new MemoryCheckpointStore(); - await store.put('rt_big', 1, Buffer.alloc(2048)); + await store.put('rt_big', 1000, 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 index 76219a2..0dff670 100644 --- a/service/src/runtime-session/checkpoint-store.ts +++ b/service/src/runtime-session/checkpoint-store.ts @@ -3,33 +3,37 @@ import { env } from '../config'; /** * Durable storage for session workspace checkpoints. Each checkpoint writes a - * distinct, strictly increasing object under a per-session prefix - * (`/.tar.gz`); restore reads the highest - * sequence. A put that timed out and lands late therefore writes an OLDER - * sequence and can never overwrite a newer checkpoint, and a relaunch can still - * find the latest checkpoint by listing the prefix even if the registry record - * was lost. Older sequences are best-effort pruned after each successful put. + * distinct object under a per-session prefix keyed by the wall-clock time the + * checkpoint was taken (`/.tar.gz`); + * restore reads the lexicographically-greatest (newest) object. A put that + * timed out and lands late carries its START time, so it writes an OLDER key + * and can never overwrite a newer checkpoint; and because the key is a + * never-resetting timestamp (not a counter that can expire), a fresh checkpoint + * after a long idle gap always sorts above and prunes any retained older + * object. 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. */ + put(runtimeSessionId: string, takenAtMs: number, data: Buffer): Promise; + /** Reads the newest 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; } export class CheckpointTooLargeError extends Error {} -/** Zero-padded so lexicographic key order matches numeric sequence order. */ -const SEQUENCE_WIDTH = 20; +/** Zero-padded so lexicographic key order matches numeric (chronological) + * order; 15 digits covers epoch-ms well past the year 30000. */ +const TIMESTAMP_WIDTH = 15; 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`; +export function checkpointObjectKey(runtimeSessionId: string, takenAtMs: number): string { + return `${checkpointPrefixFor(runtimeSessionId)}${String(takenAtMs).padStart(TIMESTAMP_WIDTH, '0')}.tar.gz`; } /** S3/MinIO-backed store using the same MINIO_* envs as file-server. */ @@ -50,8 +54,8 @@ export class MinioCheckpointStore implements CheckpointStore { 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); + async put(runtimeSessionId: string, takenAtMs: number, data: Buffer): Promise { + const key = checkpointObjectKey(runtimeSessionId, takenAtMs); await this.client.putObject(this.bucket, key, data, data.length, { 'Content-Type': 'application/x-gtar', }); @@ -119,11 +123,11 @@ export class MinioCheckpointStore implements CheckpointStore { export class MemoryCheckpointStore implements CheckpointStore { readonly objects = new Map(); - async put(runtimeSessionId: string, sequence: number, data: Buffer): Promise { - const key = checkpointObjectKey(runtimeSessionId, sequence); + async put(runtimeSessionId: string, takenAtMs: number, data: Buffer): Promise { + const key = checkpointObjectKey(runtimeSessionId, takenAtMs); const prefix = checkpointPrefixFor(runtimeSessionId); for (const existing of this.objects.keys()) { - /* prune strictly-older sequences only — never a newer one */ + /* 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)); diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts index c1c896c..ed69c63 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -4,7 +4,6 @@ import { microvmPortHeaders } from './lambda-client'; import type { CheckpointStore } from './checkpoint-store'; import { acquireRuntimeSessionLock, - allocateCheckpointSequence, readRuntimeSessionRecord, releaseRuntimeSessionLock, writeRuntimeSessionRecord, @@ -109,6 +108,12 @@ export async function checkpointSession(args: { microvmCheckpoints.inc({ outcome: 'skipped_busy' }); return 'skipped_busy'; } + /* Stamp the checkpoint's start time and key the object by it: it never + * resets (unlike a counter with a TTL), so a checkpoint after a long idle gap + * always sorts above a retained older object, and a put that stalled and + * lands late carries this earlier start time so it can't overwrite a newer + * checkpoint. */ + const takenAtMs = Date.now(); try { const record = await readRuntimeSessionRecord(args.runtimeSessionId); if (!record || record.state !== 'RUNNING' || !record.microvm_id || !record.endpoint) { @@ -121,15 +126,12 @@ export async function checkpointSession(args: { endpointBase: args.normalizeEndpoint(record.endpoint), runtimeSessionId: args.runtimeSessionId, }, args.config); - /* Each checkpoint writes a distinct, strictly increasing object key, so a - * put that stalled past the lock and lands late writes an OLDER key and can - * never overwrite the newer one restore reads. Fence the pointer write - * first: if it reports we were fenced, skip the store entirely. */ - const sequence = await allocateCheckpointSequence(args.runtimeSessionId); + /* 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(), + workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId, takenAtMs), + checkpointed_at: takenAtMs, }, lockToken); if (!persisted) { microvmCheckpoints.inc({ outcome: 'skipped_busy' }); @@ -138,7 +140,7 @@ export async function checkpointSession(args: { /* 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.store.put(args.runtimeSessionId, takenAtMs, data), args.config.timeoutMs, 'checkpoint store.put', ); diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index d57effa..7f3f23e 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -44,7 +44,6 @@ export interface RuntimeSessionRecord { 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, so @@ -202,17 +201,6 @@ export async function allocateRuntimeSessionGeneration(runtimeSessionId: string) return generation; } -/** Per-checkpoint monotonic sequence. Each successful checkpoint writes a - * distinct, strictly increasing object key so a put that timed out and lands - * late can never overwrite a newer checkpoint (restore always reads the max - * sequence). Bumps on every checkpoint, unlike the per-relaunch generation. */ -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; -} - export async function touchRuntimeSessionActive(runtimeSessionId: string, lastSeenAtMs: number): Promise { await redis.zadd(ACTIVE_ZSET, lastSeenAtMs, runtimeSessionId); } diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 32162b9..b821d67 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -14,7 +14,7 @@ import { setRedisForTests as setRegistryRedis, writeRuntimeSessionRecord, } from '../runtime-session/registry'; -import { MemoryCheckpointStore, checkpointObjectKey } from '../runtime-session/checkpoint-store'; +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'; @@ -479,18 +479,20 @@ describe('LambdaMicrovmSandboxBackend auto-checkpoint', () => { 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'); - expect(record?.workspace_checkpoint).toBe(checkpointObjectKey('rt_ckpt_1', 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', 1, Buffer.from('PRIOR_WORKSPACE')); + 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', 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); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 99c8cf3..fc78191 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -547,7 +547,21 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { } catch (error) { microvmLaunches.inc({ outcome: 'failed' }); await this.terminate(client, vm.microvmId, 'error'); - throw 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, + ); } } From 11bc26c95cf2049d90cca2c0290562285d129ac8 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 10:41:55 -0400 Subject: [PATCH 21/24] fix: address Codex review round 7 of stateful session mechanics - Re-prime session inputs without following prior-turn symlinks (P1). A persistent session workspace can contain a symlink (a file path or a parent dir) planted by earlier sandbox code; priming as root followed it and wrote/chmod'd outside the workspace. secureAncestors now rejects a symlinked ancestor (lstat, no follow), and both the inline writeFile and the download rename clear any squatted symlink/dir/file at the target first, so a prime always lands a fresh regular file inside the workspace. - Don't hide CHANGED previously-primed inputs (refines round 6). The unconditional skip suppressed a primed input even after the sandbox modified it and defeated the upload-retry path. Only suppress while the on-disk hash matches the primed baseline (or the input is read-only, whose edits are dropped by contract); a modified writable input surfaces as an output again. - Don't time out a normal auto-resume as unhealthy (P2). A suspended session VM auto-resumes on the first request, and that latency can exceed the 5s health budget, so the preflight probe misread a valid slow resume as MICROVM_UNHEALTHY and tore the VM down. Skip the preflight health check for reused session VMs (the execute carries the resume under the job budget; an evicted VM already fails token minting with not_found); freshly-launched VMs still get the probe. --- api/src/job.ts | 31 +++++++++++++++++-- api/src/session-workspace.test.ts | 4 +++ api/src/session-workspace.ts | 6 ++++ .../sandbox-backend/lambda-microvm.test.ts | 13 ++++++++ service/src/sandbox-backend/lambda-microvm.ts | 28 ++++++++++++----- 5 files changed, 73 insertions(+), 9 deletions(-) diff --git a/api/src/job.ts b/api/src/job.ts index c4243c4..0638cdc 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -799,6 +799,15 @@ 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); } @@ -1028,6 +1037,10 @@ export class Job { const finalParent = path.dirname(finalPath); 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'; @@ -1122,6 +1135,12 @@ export class Job { const parentDir = path.dirname(filePath); 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); @@ -1581,9 +1600,17 @@ export class Job { /* 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. It persists in the workspace as an input, not an output. */ + * 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)) { - return { collected: false, truncated: false, stopLoop: false }; + 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; diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index 16fd2a1..096c772 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -119,6 +119,10 @@ describe('SessionWorkspace state', () => { 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); diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index 41ce79d..04b8fac 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -163,6 +163,12 @@ export class SessionWorkspace { 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 }); } diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index b821d67..15458a9 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -330,6 +330,19 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { 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); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index fc78191..464971d 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -182,8 +182,8 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { try { const existing = await readRuntimeSessionRecord(runtimeSessionId); - const vm = await this.findOrLaunchSession(client, ctx, runtimeSessionId, existing, lockToken); - const result = await this.executeOnSessionVm(client, vm, req, ctx, runtimeSessionId, lockToken); + 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. */ @@ -238,9 +238,10 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { ctx: SandboxExecuteContext, runtimeSessionId: string, lockToken: string, + reused: boolean, ): Promise { try { - return await this.proxyExecute(client, vm, req, ctx, runtimeSessionId); + 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 @@ -269,7 +270,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { runtimeSessionId: string, record: RuntimeSessionRecord | null, lockToken: string, - ): Promise { + ): 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 @@ -295,7 +296,10 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { if (reusable && record) { /* Reuse the warm VM. If AWS auto-suspended it, the proxy request * transparently auto-resumes it (idlePolicy.autoResume). */ - return { microvmId: record.microvm_id as string, state: 'RUNNING', endpoint: record.endpoint }; + 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 @@ -368,7 +372,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { config: this.config.checkpoint, }); } - return vm; + return { vm, reused: false }; } /** @@ -452,10 +456,20 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { req: SandboxTransportRequest, ctx: SandboxExecuteContext, runtimeSessionId?: string, + reused = false, ): Promise { const base = normalizeMicrovmEndpoint(vm.endpoint ?? ''); const token = await this.mintAuthToken(client, vm.microvmId); - await this.assertHealthy(base, token.token, ctx); + /* 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 From 61ba32e852d3e534dac129de5f167523ba79444a Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 10:49:44 -0400 Subject: [PATCH 22/24] docs(lambda-microvm): guard the build-log path against the wrong "fix" MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AWS docs list the MicroVM build log group as /aws/lambda/microvms/ (slash), but the real path is /aws/lambda-microvms/ (hyphen) — re-verified against a live account (every build's group is the hyphen form; the slash path does not exist). A passive "not the docs' path" note wasn't enough (a reviewer still flagged it citing the docs), so make it an explicit guard in both the terraform comment and the runbook: do not correct it to the slash path or the build logs are lost and CreateMicrovmImage fails with an empty stateReason. --- docs/lambda-microvm/README.md | 7 +++++-- docs/lambda-microvm/terraform/main.tf | 14 +++++++++----- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md index 475c357..c57ee15 100644 --- a/docs/lambda-microvm/README.md +++ b/docs/lambda-microvm/README.md @@ -304,8 +304,11 @@ Each of these cost a silent or blind failure during bring-up: 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 - docs' `/aws/lambda/microvms/`. +- **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`. diff --git a/docs/lambda-microvm/terraform/main.tf b/docs/lambda-microvm/terraform/main.tf index 5e45fd9..34abdf6 100644 --- a/docs/lambda-microvm/terraform/main.tf +++ b/docs/lambda-microvm/terraform/main.tf @@ -103,11 +103,15 @@ resource "aws_s3_bucket_lifecycle_configuration" "checkpoint" { # -------------------------------------------------------------------------- # CloudWatch Logs -# Build logs land at the EXACT path `/aws/lambda-microvms/` (hyphen, -# not the docs' `/aws/lambda/microvms/...`). 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. +# 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}" From a53ff5f5f65a9aad83b6a826fe840261eb35d0d7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 11:18:50 -0400 Subject: [PATCH 23/24] fix: address Codex review round 8 of stateful session mechanics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Advance last_seen_at on checkpointed executes. checkpointUnderLock returned the record checkpointSession re-read (built from the pre-execute snapshot), so liveness never moved forward and an actively-used session eventually looked idle and relaunched needlessly. Re-apply last_seen_at: now. - Budget the session lock for all checkpoint I/O. Restore (store.get + pushRestore) and the post-run checkpoint (pullCheckpoint + store.put) each do TWO timeout-bounded I/Os, so the TTL now reserves 4x CHECKPOINT_TIMEOUT (was 2x) or the lock could expire mid-op and let a second worker run concurrently. - Include network-connector drift in session reuse (P1). configMatches only compared image/version/port; connectors apply only at RunMicrovm, so a hardened deploy that tightened egress kept reusing VMs on the old broader policy. Record a connector fingerprint at launch and compare it. - Wait for runner health before restoring. Restore ran as soon as the VM was RUNNING (control-plane state) — before the app listener was up — so pushRestore could race the boot, fail non-fatally, and run the first execute on an empty workspace. Poll readiness first; tear the VM down if it never comes up. - Treat an empty CODEAPI_CHECKPOINT_BUCKET as unset (|| not ??) so it falls through to MINIO_BUCKET instead of selecting '' and failing every S3 op. - Don't clobber a user file at the checkpoint sidecar path. If user code creates a real .codeapi-session-meta.json, skip metadata persistence that turn rather than deleting it, and on restore only remove a sidecar we recognize as ours. --- api/src/session-checkpoint.ts | 35 ++++++---- .../src/runtime-session/checkpoint-store.ts | 5 +- service/src/runtime-session/registry.ts | 15 +++-- .../sandbox-backend/lambda-microvm.test.ts | 14 ++++ service/src/sandbox-backend/lambda-microvm.ts | 67 +++++++++++++++++-- 5 files changed, 113 insertions(+), 23 deletions(-) diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index 7f67e45..018cbe9 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -36,17 +36,25 @@ export async function streamSessionCheckpoint(res: Response): Promise { /* 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 from the live - * workspace once tar has read it. Sandboxed code from a prior exec can squat - * this name as a symlink OR a directory; the API process runs as root, so - * remove whatever squats the path first — `recursive` clears a directory - * (else the write later fails EISDIR and breaks every checkpoint), `force` - * ignores absence, and neither follows a symlink — then create a fresh - * regular file exclusively (`wx`), so a privileged write can't follow a link - * and clobber an arbitrary target outside the workspace. */ + * 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); - await fsp.rm(metaPath, { force: true, recursive: true }); - await fsp.writeFile(metaPath, JSON.stringify(session.snapshotMeta()), { flag: 'wx' }); + 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(session.snapshotMeta()), { flag: 'wx' }); + wroteSidecar = true; + } res.status(200); res.setHeader('Content-Type', CHECKPOINT_CONTENT_TYPE); @@ -63,7 +71,7 @@ export async function streamSessionCheckpoint(res: Response): Promise { if (!res.headersSent) res.status(500).json({ message: 'checkpoint failed' }); else res.destroy(); } finally { - await fsp.rm(metaPath, { force: true }).catch(() => {}); + if (wroteSidecar) await fsp.rm(metaPath, { force: true }).catch(() => {}); } } @@ -125,11 +133,12 @@ async function applyRestoredMeta(session: SessionWorkspace, dir: string): Promis const parsed = JSON.parse(await fsp.readFile(metaPath, 'utf8')) as SessionMetaSnapshot; if (Array.isArray(parsed?.primed) && Array.isArray(parsed?.surfaced)) { session.loadMeta(parsed); + /* Only remove a sidecar we recognize as our own metadata — a user file + * that merely shares the reserved name is left in the workspace intact. */ + await fsp.rm(metaPath, { force: true }).catch(() => {}); } } catch (error) { logger.debug({ err: error }, 'No session meta sidecar to restore'); - } finally { - await fsp.rm(metaPath, { force: true }).catch(() => {}); } } diff --git a/service/src/runtime-session/checkpoint-store.ts b/service/src/runtime-session/checkpoint-store.ts index 0dff670..cfb5402 100644 --- a/service/src/runtime-session/checkpoint-store.ts +++ b/service/src/runtime-session/checkpoint-store.ts @@ -51,7 +51,10 @@ export class MinioCheckpointStore implements CheckpointStore { ...(process.env.MINIO_SESSION_TOKEN ? { sessionToken: process.env.MINIO_SESSION_TOKEN } : {}), ...(process.env.MINIO_REGION ? { region: process.env.MINIO_REGION } : {}), }); - this.bucket = process.env.CODEAPI_CHECKPOINT_BUCKET ?? process.env.MINIO_BUCKET ?? 'test-bucket'; + /* `||` 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, takenAtMs: number, data: Buffer): Promise { diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index 7f3f23e..2fd0748 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -31,6 +31,11 @@ export interface RuntimeSessionRecord { 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; @@ -52,16 +57,18 @@ const ACTIVE_ZSET = 'rtsx:active'; * session concurrently. A relaunch path can spend, back to back: * - launch throttle wait (acquireOpBudget) ≤ LAUNCH_TIMEOUT_MS * - poll to RUNNING (waitUntilRunning) ≤ LAUNCH_TIMEOUT_MS - * - restore the checkpoint before the first exec ≤ CHECKPOINT_TIMEOUT_MS + * - restore: store.get + pushRestore ≤ 2× CHECKPOINT_TIMEOUT_MS * - health check ≤ HEALTH_TIMEOUT_MS * - the execute ≤ JOB_TIMEOUT - * - the post-run checkpoint ≤ CHECKPOINT_TIMEOUT_MS - * so budget 2× launch and 2× checkpoint, plus headroom for lock-wait + jitter. */ + * - post-run checkpoint: pullCheckpoint + store.put ≤ 2× CHECKPOINT_TIMEOUT_MS + * Restore AND the post-run checkpoint each do TWO timeout-bounded I/Os (an + * object-store call plus a runner proxy call), so budget 2× launch and 4× + * checkpoint, plus headroom for lock-wait + jitter. */ export const RUNTIME_SESSION_LOCK_TTL_MS = env.JOB_TIMEOUT + 2 * env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + env.LAMBDA_MICROVM_HEALTH_TIMEOUT_MS + - 2 * env.CHECKPOINT_TIMEOUT_MS + + 4 * env.CHECKPOINT_TIMEOUT_MS + 60_000; const MAX_MICROVM_DURATION_SECONDS = 28_800; diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index 15458a9..d17197e 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -464,6 +464,20 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { 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', () => { diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 464971d..86e922c 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -264,6 +264,15 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { } } + /** 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, @@ -279,7 +288,8 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { const configMatches = record && record.image_arn === this.config.imageArn && record.image_version === this.config.imageVersion - && record.port === this.config.port; + && 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 @@ -345,6 +355,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { port: this.config.port, image_arn: this.config.imageArn, image_version: this.config.imageVersion, + connectors: this.connectorFingerprint(), state: 'RUNNING', generation, launched_at: launchedAt, @@ -363,18 +374,60 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { * 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; + } await restoreSession({ mintToken: (microvmId) => this.mintAuthToken(client, microvmId), store: this.checkpointStore, runtimeSessionId, microvmId: vm.microvmId, - endpointBase: normalizeMicrovmEndpoint(vm.endpoint ?? ''), + endpointBase, config: this.config.checkpoint, }); } 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) @@ -400,10 +453,14 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { normalizeEndpoint: normalizeMicrovmEndpoint, lockToken, }); - /* checkpointSession wrote the pointer under our lock on success; re-read - * so we don't clobber it with our stale `base`. */ + /* 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') { - return (await readRuntimeSessionRecord(runtimeSessionId)) ?? base; + const persisted = await readRuntimeSessionRecord(runtimeSessionId); + return persisted ? { ...persisted, last_seen_at: now } : base; } return base; } From 78712876878ba6360eda9c2748504c3838ae2f8c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 12:03:00 -0400 Subject: [PATCH 24/24] fix: address Codex review round 9 of stateful session mechanics - Renew the session lock on a heartbeat instead of pinning its TTL to the worst-case sum. The critical path (launch throttle + readiness/restore + execute + checkpoint) includes several per-op token-mint throttle waits, so no fixed TTL reliably covers it (this is the third TTL finding). renewRuntime- SessionLock extends the lock while the holder runs; a fenced renew stops when another worker owns it. TTL is now just a comfortable base. - Key checkpoints by a store-seeded monotonic Redis sequence, not wall-clock. Wall-clock keys (round 6) mis-order across pods with skewed/backward clocks; the sequence is a pure INCR seeded above latestSequence() after a TTL reset, so it's skew-free AND never resets below retained objects. - Recycle a session VM when push-restore fails. restoreSession now distinguishes fetch_failed (clean fresh workspace, safe to execute) from push_failed (runner may be mid-extract/wipe) and tears the VM down on push_failed so the next call relaunches clean. - Recycle a reused VM on a proxy 502/503/504 (a suspended VM that failed to auto-resume) instead of treating it as a live-runner response and reusing the dead VM until idle expiry. - Authenticity marker on the checkpoint sidecar: restore only loads/deletes a file carrying our marker, so a user file that merely shares the reserved name and happens to hold primed/surfaced arrays is never poisoned/deleted. - Track inline source files (writeFile) as session inputs so switching entrypoint (main.py -> script.sh) doesn't re-surface the prior source as a generated file. - Build session workspace ancestors no-follow BEFORE mkdir. The round-7 symlink check ran after mkdir -p, which had already followed a symlinked ancestor; ensureDirNoFollow creates one component at a time, rejecting symlinks. --- api/src/job.ts | 48 ++++++++++++- api/src/session-checkpoint.ts | 12 ++-- api/src/session-workspace.ts | 7 ++ .../runtime-session/checkpoint-store.test.ts | 47 ++++++------ .../src/runtime-session/checkpoint-store.ts | 71 ++++++++++++------- service/src/runtime-session/checkpoint.ts | 42 +++++++---- service/src/runtime-session/registry.ts | 70 ++++++++++++++---- .../sandbox-backend/lambda-microvm.test.ts | 13 ++++ service/src/sandbox-backend/lambda-microvm.ts | 38 +++++++++- 9 files changed, 268 insertions(+), 80 deletions(-) diff --git a/api/src/job.ts b/api/src/job.ts index 0638cdc..ad4c5b7 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -788,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; @@ -1035,7 +1071,10 @@ 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 @@ -1133,7 +1172,8 @@ 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 @@ -1146,6 +1186,10 @@ export class Job { 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( diff --git a/api/src/session-checkpoint.ts b/api/src/session-checkpoint.ts index 018cbe9..f0be84d 100644 --- a/api/src/session-checkpoint.ts +++ b/api/src/session-checkpoint.ts @@ -6,7 +6,7 @@ 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, getBoundSessionWorkspace } from './session-workspace'; +import { SESSION_META_FILE, SESSION_META_MARKER, getBoundSessionWorkspace } from './session-workspace'; /** * Session workspace checkpoint / restore. @@ -52,7 +52,7 @@ export async function streamSessionCheckpoint(res: Response): Promise { 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(session.snapshotMeta()), { flag: 'wx' }); + await fsp.writeFile(metaPath, JSON.stringify({ marker: SESSION_META_MARKER, ...session.snapshotMeta() }), { flag: 'wx' }); wroteSidecar = true; } @@ -131,10 +131,12 @@ async function applyRestoredMeta(session: SessionWorkspace, dir: string): Promis const stat = await fsp.lstat(metaPath).catch(() => null); if (!stat?.isFile()) return; const parsed = JSON.parse(await fsp.readFile(metaPath, 'utf8')) as SessionMetaSnapshot; - if (Array.isArray(parsed?.primed) && Array.isArray(parsed?.surfaced)) { + /* 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); - /* Only remove a sidecar we recognize as our own metadata — a user file - * that merely shares the reserved name is left in the workspace intact. */ await fsp.rm(metaPath, { force: true }).catch(() => {}); } } catch (error) { diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index 04b8fac..fa0dce6 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -42,7 +42,14 @@ export const RUNTIME_SESSION_ID_HEADER = 'x-runtime-session-id'; * 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]>; } diff --git a/service/src/runtime-session/checkpoint-store.test.ts b/service/src/runtime-session/checkpoint-store.test.ts index e321df7..3554419 100644 --- a/service/src/runtime-session/checkpoint-store.test.ts +++ b/service/src/runtime-session/checkpoint-store.test.ts @@ -9,20 +9,19 @@ import { const BIG = 1_000_000; describe('checkpoint store', () => { - test('object key is per session + timestamp, zero-padded for lexical order', () => { - expect(checkpointObjectKey('rt_abc', 1)).toBe('rtsx-checkpoints/rt_abc/000000000000001.tar.gz'); + 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 == chronological order across widths */ + /* 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_abc', 1_783_000_000_000) > checkpointObjectKey('rt_abc', 999)).toBe(true); expect(checkpointObjectKey('rt_xyz', 1)).not.toBe(checkpointObjectKey('rt_abc', 1)); }); - test('memory store round-trips the newest bytes and copies defensively', async () => { + 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', 1000, original); + await store.put('rt_1', 1, original); const fetched = await store.get('rt_1', BIG); expect(fetched?.toString()).toBe('workspace-bytes'); @@ -36,30 +35,38 @@ describe('checkpoint store', () => { expect(await store.get('rt_missing', BIG)).toBeNull(); }); - test('get reads the newest timestamp, not the last write', async () => { + test('get reads the highest sequence, not the last write', async () => { const store = new MemoryCheckpointStore(); - await store.put('rt_1', 2000, Buffer.from('newer')); - /* a stale put carrying an EARLIER start time lands late but must NOT win — - * and a fresh checkpoint after an idle gap (later timestamp) always does */ - await store.put('rt_1', 1000, Buffer.from('stale')); + 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('put prunes older timestamps for the session', async () => { + test('latestSequence reports the max retained sequence (for reset seeding)', async () => { const store = new MemoryCheckpointStore(); - await store.put('rt_1', 1000, Buffer.from('v1')); - await store.put('rt_1', 2000, Buffer.from('v2')); - /* only the newest object survives; siblings for other sessions are untouched */ - await store.put('rt_other', 1000, Buffer.from('other')); + 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', 2000)); - expect(keys).not.toContain(checkpointObjectKey('rt_1', 1000)); - expect(keys).toContain(checkpointObjectKey('rt_other', 1000)); + 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', 1000, Buffer.alloc(2048)); + 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 index cfb5402..957c7ac 100644 --- a/service/src/runtime-session/checkpoint-store.ts +++ b/service/src/runtime-session/checkpoint-store.ts @@ -3,37 +3,46 @@ import { env } from '../config'; /** * Durable storage for session workspace checkpoints. Each checkpoint writes a - * distinct object under a per-session prefix keyed by the wall-clock time the - * checkpoint was taken (`/.tar.gz`); - * restore reads the lexicographically-greatest (newest) object. A put that - * timed out and lands late carries its START time, so it writes an OLDER key - * and can never overwrite a newer checkpoint; and because the key is a - * never-resetting timestamp (not a counter that can expire), a fresh checkpoint - * after a long idle gap always sorts above and prunes any retained older - * object. A relaunch finds the latest by listing the prefix even with no + * 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, takenAtMs: number, data: Buffer): Promise; - /** Reads the newest 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. */ + 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 (chronological) - * order; 15 digits covers epoch-ms well past the year 30000. */ -const TIMESTAMP_WIDTH = 15; +/** 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, takenAtMs: number): string { - return `${checkpointPrefixFor(runtimeSessionId)}${String(takenAtMs).padStart(TIMESTAMP_WIDTH, '0')}.tar.gz`; +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. */ @@ -57,17 +66,22 @@ export class MinioCheckpointStore implements CheckpointStore { this.bucket = process.env.CODEAPI_CHECKPOINT_BUCKET || process.env.MINIO_BUCKET || 'test-bucket'; } - async put(runtimeSessionId: string, takenAtMs: number, data: Buffer): Promise { - const key = checkpointObjectKey(runtimeSessionId, takenAtMs); + 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 stale late put must not delete the checkpoint that superseded it — - * restore always reads the max sequence). */ + * (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; @@ -126,8 +140,8 @@ export class MinioCheckpointStore implements CheckpointStore { export class MemoryCheckpointStore implements CheckpointStore { readonly objects = new Map(); - async put(runtimeSessionId: string, takenAtMs: number, data: Buffer): Promise { - const key = checkpointObjectKey(runtimeSessionId, takenAtMs); + 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 */ @@ -136,6 +150,15 @@ export class MemoryCheckpointStore implements CheckpointStore { 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; diff --git a/service/src/runtime-session/checkpoint.ts b/service/src/runtime-session/checkpoint.ts index ed69c63..afb2351 100644 --- a/service/src/runtime-session/checkpoint.ts +++ b/service/src/runtime-session/checkpoint.ts @@ -4,8 +4,10 @@ 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'; @@ -108,12 +110,6 @@ export async function checkpointSession(args: { microvmCheckpoints.inc({ outcome: 'skipped_busy' }); return 'skipped_busy'; } - /* Stamp the checkpoint's start time and key the object by it: it never - * resets (unlike a counter with a TTL), so a checkpoint after a long idle gap - * always sorts above a retained older object, and a put that stalled and - * lands late carries this earlier start time so it can't overwrite a newer - * checkpoint. */ - const takenAtMs = Date.now(); try { const record = await readRuntimeSessionRecord(args.runtimeSessionId); if (!record || record.state !== 'RUNNING' || !record.microvm_id || !record.endpoint) { @@ -126,12 +122,25 @@ export async function checkpointSession(args: { 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, takenAtMs), - checkpointed_at: takenAtMs, + workspace_checkpoint: checkpointObjectKey(args.runtimeSessionId, sequence), + checkpointed_at: Date.now(), }, lockToken); if (!persisted) { microvmCheckpoints.inc({ outcome: 'skipped_busy' }); @@ -140,7 +149,7 @@ export async function checkpointSession(args: { /* 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, takenAtMs, data), + args.store.put(args.runtimeSessionId, sequence, data), args.config.timeoutMs, 'checkpoint store.put', ); @@ -168,7 +177,7 @@ export async function restoreSession(args: { microvmId: string; endpointBase: string; config: CheckpointConfig; -}): Promise<'restored' | 'absent' | 'failed'> { +}): 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 @@ -181,12 +190,15 @@ export async function restoreSession(args: { '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 'failed'; + return 'fetch_failed'; } if (!data) { microvmRestores.inc({ outcome: 'absent' }); @@ -205,11 +217,15 @@ export async function restoreSession(args: { }); 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 restore failed; continuing with a fresh workspace', { + logger.warn('Checkpoint push-restore failed; the VM workspace may be partial', { runtimeSessionId: args.runtimeSessionId, error: error instanceof Error ? error.message : String(error), }); - return 'failed'; + return 'push_failed'; } } diff --git a/service/src/runtime-session/registry.ts b/service/src/runtime-session/registry.ts index 2fd0748..578f084 100644 --- a/service/src/runtime-session/registry.ts +++ b/service/src/runtime-session/registry.ts @@ -49,21 +49,18 @@ export interface RuntimeSessionRecord { 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, so - * its TTL must outlive the worst-case sum of the configurable budgets — else a - * live holder is fenced mid-operation and a second worker acquires the same - * session concurrently. A relaunch path can spend, back to back: - * - launch throttle wait (acquireOpBudget) ≤ LAUNCH_TIMEOUT_MS - * - poll to RUNNING (waitUntilRunning) ≤ LAUNCH_TIMEOUT_MS - * - restore: store.get + pushRestore ≤ 2× CHECKPOINT_TIMEOUT_MS - * - health check ≤ HEALTH_TIMEOUT_MS - * - the execute ≤ JOB_TIMEOUT - * - post-run checkpoint: pullCheckpoint + store.put ≤ 2× CHECKPOINT_TIMEOUT_MS - * Restore AND the post-run checkpoint each do TWO timeout-bounded I/Os (an - * object-store call plus a runner proxy call), so budget 2× launch and 4× - * checkpoint, plus headroom for lock-wait + jitter. */ +/** 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 + @@ -76,6 +73,7 @@ export const RUNTIME_SESSION_RECORD_TTL_SECONDS = MAX_MICROVM_DURATION_SECONDS + type RedisWithScripts = Redis & { releaseRuntimeSessionLockScript(lockKey: string, token: string): Promise; + renewRuntimeSessionLockScript(lockKey: string, token: string, ttlMs: string): Promise; writeRuntimeSessionRecordScript( sessKey: string, lockKey: string, @@ -101,6 +99,10 @@ function registerScripts(client: Redis): RedisWithScripts { 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 @@ -168,6 +170,29 @@ export async function releaseRuntimeSessionLock(runtimeSessionId: string, token: } } +/** 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; @@ -208,6 +233,25 @@ export async function allocateRuntimeSessionGeneration(runtimeSessionId: string) 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); } diff --git a/service/src/sandbox-backend/lambda-microvm.test.ts b/service/src/sandbox-backend/lambda-microvm.test.ts index d17197e..9b2888a 100644 --- a/service/src/sandbox-backend/lambda-microvm.test.ts +++ b/service/src/sandbox-backend/lambda-microvm.test.ts @@ -357,6 +357,19 @@ describe('LambdaMicrovmSandboxBackend session execution', () => { 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); diff --git a/service/src/sandbox-backend/lambda-microvm.ts b/service/src/sandbox-backend/lambda-microvm.ts index 86e922c..a10a9bc 100644 --- a/service/src/sandbox-backend/lambda-microvm.ts +++ b/service/src/sandbox-backend/lambda-microvm.ts @@ -9,10 +9,12 @@ import { LambdaMicrovmApiError, microvmPortHeaders } from '../runtime-session/la 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, @@ -180,6 +182,18 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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); @@ -214,6 +228,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { await touchRuntimeSessionActive(runtimeSessionId, now); return result; } finally { + clearInterval(heartbeat); await releaseRuntimeSessionLock(runtimeSessionId, lockToken); } } @@ -253,10 +268,17 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { * 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. */ + * (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'; - if (ctx.signal.aborted || transportFailure || 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(() => {}); } @@ -388,7 +410,7 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { await removeRuntimeSession(runtimeSessionId, lockToken).catch(() => {}); throw error; } - await restoreSession({ + const restoreResult = await restoreSession({ mintToken: (microvmId) => this.mintAuthToken(client, microvmId), store: this.checkpointStore, runtimeSessionId, @@ -396,6 +418,16 @@ export class LambdaMicrovmSandboxBackend implements SandboxBackend { 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 }; }