From 8a40cc50b05a48056194fb52023a3ea59adb76af Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 14 Aug 2026 12:07:30 -0400 Subject: [PATCH] feat: isolate default and stateful execution profiles --- README.md | 21 ++++ docs/lambda-microvm/README.md | 2 + helm/codeapi/README.md | 9 ++ helm/codeapi/templates/api-deployment.yaml | 2 + .../templates/worker-sandbox-deployment.yaml | 2 + helm/codeapi/values.yaml | 5 + service/openapi.yml | 107 +++++++++++++++++- service/src/api-server.ts | 2 + service/src/config.ts | 18 ++- service/src/enum/service.ts | 5 - service/src/execution-profile.test.ts | 83 ++++++++++++++ service/src/execution-profile.ts | 97 ++++++++++++++++ service/src/lifecycle.ts | 10 +- service/src/local-api.ts | 4 + service/src/metrics.ts | 19 ++++ .../src/middleware/execution-profile.test.ts | 78 +++++++++++++ service/src/middleware/execution-profile.ts | 37 ++++++ service/src/queue.ts | 20 ++-- service/src/secure-startup.test.ts | 58 ++++++++++ service/src/secure-startup.ts | 35 ++++++ service/src/service-api.ts | 2 + service/src/service/router.ts | 6 +- service/src/workers.ts | 7 +- 23 files changed, 606 insertions(+), 23 deletions(-) create mode 100644 service/src/execution-profile.test.ts create mode 100644 service/src/execution-profile.ts create mode 100644 service/src/middleware/execution-profile.test.ts create mode 100644 service/src/middleware/execution-profile.ts diff --git a/README.md b/README.md index 718db6a..f848b1d 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,27 @@ Code Interpreter (internally `codeapi`, the prefix used by its env vars, images, 4. Files are persisted/retrieved via the **File Server** (backed by S3) 5. Tool calls from within sandboxes are routed through the **Tool Call Server** +## Execution profiles + +Code API can run two isolated deployments at the same time: + +- `default`: the AWS-free HTTP/libkrun path, with stateless executions. +- `stateful`: the AWS Lambda MicroVM path, with runtime-session affinity. + +Set `CODEAPI_EXECUTION_PROFILE` consistently on an API deployment and its +workers. The default profile keeps the existing `python-queue` and +`other-queue`; the stateful profile uses `stateful-python-queue` and +`stateful-other-queue`. This allows both deployments to share Redis without +cross-consuming jobs. + +Trusted callers should send `X-CodeAPI-Expected-Profile: default|stateful` on +every Code API request. A request that reaches the wrong deployment fails +before enqueue with HTTP 409 and `code=execution_profile_mismatch`; every +response advertises the actual deployment in `X-CodeAPI-Execution-Profile`. +Omitting the expected-profile header remains supported for older clients, but +provides no wrong-endpoint protection. There is deliberately no silent +fallback between profiles and no automatic workspace or file migration. + ## Sandbox Isolation Two modes are supported: diff --git a/docs/lambda-microvm/README.md b/docs/lambda-microvm/README.md index adbad83..4bd3e87 100644 --- a/docs/lambda-microvm/README.md +++ b/docs/lambda-microvm/README.md @@ -246,6 +246,7 @@ builds. ```bash CODEAPI_SANDBOX_BACKEND=lambda-microvm +CODEAPI_EXECUTION_PROFILE=stateful CODEAPI_RUNTIME_SESSION_MODE=affinity # warm sessions + checkpoints LAMBDA_MICROVM_IMAGE_ARN= LAMBDA_MICROVM_IMAGE_VERSION= # required for affinity/strict @@ -310,6 +311,7 @@ appear in `api/src/config.ts`. | Env | Default | Meaning | |---|---|---| | `CODEAPI_SANDBOX_BACKEND` | `http` | `http` (byte-identical to today) or `lambda-microvm`. | +| `CODEAPI_EXECUTION_PROFILE` | inferred | `default` for the HTTP/stateless deployment or `stateful` for the Lambda affinity/strict deployment. Stateful API and worker processes consume isolated BullMQ queues. | | `CODEAPI_RUNTIME_SESSION_MODE` | `stateless` | `stateless` \| `affinity` \| `strict`. `affinity` and `strict` require the `lambda-microvm` backend. See [Operating modes](#operating-modes). | | `CODEAPI_RUNTIME_SESSION_LOCK_WAIT_MS` | `15000` | How long a stateful execution waits for the session lock before returning `RUNTIME_SESSION_BUSY` (HTTP 409). | diff --git a/helm/codeapi/README.md b/helm/codeapi/README.md index 9198b1a..83cf2a2 100644 --- a/helm/codeapi/README.md +++ b/helm/codeapi/README.md @@ -53,6 +53,15 @@ platform rather than templated here: external ingress/service mesh, KEDA-style queue-depth autoscaling, and cloud-IAM secret delivery (the env hooks below cover all of them). +**Execution profile.** This chart deploys the AWS-free `default` profile and +sets `CODEAPI_EXECUTION_PROFILE=default` on both API and worker pods. That +profile requires the HTTP sandbox backend in stateless mode and retains the +existing `python-queue` / `other-queue` BullMQ names. A separate stateful +Lambda MicroVM deployment must use `CODEAPI_EXECUTION_PROFILE=stateful`; it +then consumes `stateful-python-queue` / `stateful-other-queue`, so both stacks +may safely share Redis without consuming each other's jobs. Do not mix API +and worker profile values within one deployment. + **Authentication.** Outside local mode the API verifies JWTs. Configure the verifier through environment variables on the api component, e.g.: diff --git a/helm/codeapi/templates/api-deployment.yaml b/helm/codeapi/templates/api-deployment.yaml index 0c801ab..7a62167 100644 --- a/helm/codeapi/templates/api-deployment.yaml +++ b/helm/codeapi/templates/api-deployment.yaml @@ -41,6 +41,8 @@ spec: {{ include "codeapi.otel.env" (dict "root" . "serviceName" "aiml-codeapi-api") | nindent 12 }} - name: CODEAPI_HARDENED_SANDBOX_MODE value: {{ .Values.hardenedSandboxMode | quote }} + - name: CODEAPI_EXECUTION_PROFILE + value: {{ .Values.executionProfile | quote }} # Redis connection - name: REDIS_HOST value: {{ include "codeapi.redis.host" . }} diff --git a/helm/codeapi/templates/worker-sandbox-deployment.yaml b/helm/codeapi/templates/worker-sandbox-deployment.yaml index 3484a61..c1bea5e 100644 --- a/helm/codeapi/templates/worker-sandbox-deployment.yaml +++ b/helm/codeapi/templates/worker-sandbox-deployment.yaml @@ -147,6 +147,8 @@ spec: {{ include "codeapi.otel.env" (dict "root" . "serviceName" "aiml-codeapi-service-worker") | nindent 12 }} - name: CODEAPI_HARDENED_SANDBOX_MODE value: {{ .Values.hardenedSandboxMode | quote }} + - name: CODEAPI_EXECUTION_PROFILE + value: {{ .Values.executionProfile | quote }} - name: SANDBOX_ENDPOINT value: "http://{{ include "codeapi.fullname" . }}-sandbox-runner:{{ .Values.workerSandbox.sandbox.port }}/api/v2" - name: EGRESS_GATEWAY_URL diff --git a/helm/codeapi/values.yaml b/helm/codeapi/values.yaml index 2385814..0f9a225 100644 --- a/helm/codeapi/values.yaml +++ b/helm/codeapi/values.yaml @@ -22,6 +22,11 @@ internalServiceAuth: hardenedSandboxMode: true +# Stable identity advertised by this API/worker deployment. The bundled chart +# is the AWS-free HTTP/libkrun profile. A separate Lambda MicroVM deployment +# must set this to `stateful`; the service then uses isolated BullMQ queues. +executionProfile: default + otel: enabled: false # OTLP/HTTP collector endpoint, e.g. "http://opentelemetry-collector.observability:4318". diff --git a/service/openapi.yml b/service/openapi.yml index c1f8f6e..16d36d7 100644 --- a/service/openapi.yml +++ b/service/openapi.yml @@ -2,7 +2,10 @@ openapi: '3.0.0' info: title: LibreChat Code Interpreter API version: '1.0.0' - description: API for sandbox code execution and file management + description: >- + API for sandbox code execution and file management. Trusted callers should + assert the intended deployment with X-CodeAPI-Expected-Profile on every + request; responses advertise the actual profile. servers: - url: https://api.librechat.ai/v1 description: LibreChat API server @@ -17,6 +20,45 @@ components: scheme: bearer bearerFormat: JWT + parameters: + ExpectedExecutionProfile: + name: X-CodeAPI-Expected-Profile + in: header + required: false + description: >- + Trusted routing assertion. A mismatched endpoint returns HTTP 409 + before any work is enqueued. Optional only for backwards compatibility. + schema: + type: string + enum: [default, stateful] + + headers: + ExecutionProfile: + description: Execution profile served by this deployment. + schema: + type: string + enum: [default, stateful] + + responses: + InvalidExecutionProfile: + description: Invalid expected execution profile + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' + content: + application/json: + schema: + $ref: '#/components/schemas/ExecutionProfileError' + ExecutionProfileMismatch: + description: The request reached a different execution profile + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' + content: + application/json: + schema: + $ref: '#/components/schemas/ExecutionProfileError' + schemas: FileRef: type: object @@ -108,6 +150,12 @@ components: type: array items: $ref: '#/components/schemas/RequestFile' + runtime_session_hint: + type: string + description: >- + Stable opaque hint for stateful runtime reuse. The server binds it + to the authenticated tenant and user. Ignored by the default + stateless profile. FileObject: type: object @@ -156,12 +204,29 @@ components: details: type: string + ExecutionProfileError: + type: object + required: [error, code, actual_profile] + properties: + error: + type: string + code: + type: string + enum: [invalid_execution_profile, execution_profile_mismatch] + expected_profile: + type: string + actual_profile: + type: string + enum: [default, stateful] + paths: /exec: post: summary: Execute code description: Execute code with specified language and parameters operationId: executeCode + parameters: + - $ref: '#/components/parameters/ExpectedExecutionProfile' requestBody: required: true content: @@ -171,6 +236,9 @@ paths: responses: '200': description: Successful execution + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' content: application/json: schema: @@ -181,6 +249,10 @@ paths: application/json: schema: $ref: '#/components/schemas/Error' + '400': + $ref: '#/components/responses/InvalidExecutionProfile' + '409': + $ref: '#/components/responses/ExecutionProfileMismatch' '503': description: Service unavailable content: @@ -192,6 +264,7 @@ paths: get: summary: Download a file parameters: + - $ref: '#/components/parameters/ExpectedExecutionProfile' - name: session_id in: path required: true @@ -205,6 +278,9 @@ paths: responses: '200': description: File content + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' content: application/octet-stream: schema: @@ -216,10 +292,16 @@ paths: application/json: schema: $ref: '#/components/schemas/Error' + '400': + $ref: '#/components/responses/InvalidExecutionProfile' + '409': + $ref: '#/components/responses/ExecutionProfileMismatch' /upload: post: summary: Upload files + parameters: + - $ref: '#/components/parameters/ExpectedExecutionProfile' requestBody: required: true content: @@ -237,6 +319,9 @@ paths: responses: '200': description: Successful upload + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' content: application/json: schema: @@ -247,11 +332,16 @@ paths: application/json: schema: $ref: '#/components/schemas/Error' + '400': + $ref: '#/components/responses/InvalidExecutionProfile' + '409': + $ref: '#/components/responses/ExecutionProfileMismatch' /files/{session_id}: get: summary: Get files information parameters: + - $ref: '#/components/parameters/ExpectedExecutionProfile' - name: session_id in: path required: true @@ -265,17 +355,25 @@ paths: responses: '200': description: Files information + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' content: application/json: schema: type: array items: $ref: '#/components/schemas/FileObject' + '400': + $ref: '#/components/responses/InvalidExecutionProfile' + '409': + $ref: '#/components/responses/ExecutionProfileMismatch' /files/{session_id}/{fileId}: delete: summary: Delete a file parameters: + - $ref: '#/components/parameters/ExpectedExecutionProfile' - name: session_id in: path required: true @@ -289,9 +387,16 @@ paths: responses: '200': description: File deleted successfully + headers: + X-CodeAPI-Execution-Profile: + $ref: '#/components/headers/ExecutionProfile' '500': description: Error deleting file content: application/json: schema: $ref: '#/components/schemas/Error' + '400': + $ref: '#/components/responses/InvalidExecutionProfile' + '409': + $ref: '#/components/responses/ExecutionProfileMismatch' diff --git a/service/src/api-server.ts b/service/src/api-server.ts index 89aba48..7868982 100644 --- a/service/src/api-server.ts +++ b/service/src/api-server.ts @@ -21,6 +21,7 @@ import programmaticRouter from './service/programmatic-router'; import { connection } from './queue'; import { metricsHandler } from './metrics'; import { httpMetricsMiddleware } from './middleware/httpMetrics'; +import { executionProfileMiddleware } from './middleware/execution-profile'; import { traceHttpRequest } from './telemetry'; import { env } from './config'; import logger from './logger'; @@ -32,6 +33,7 @@ app.disable('x-powered-by'); app.set('trust proxy', 1); app.use(traceHttpRequest('codeapi.api.request')); app.use(httpMetricsMiddleware); +app.use(executionProfileMiddleware); const v1 = Router(); diff --git a/service/src/config.ts b/service/src/config.ts index c9dcb68..4309af1 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -3,6 +3,7 @@ dotenv.config(); import { nanoid } from 'nanoid'; import type * as t from './types'; import { Languages } from './enum'; +import { resolveExecutionProfile } from './execution-profile'; export const languageConfig: Record = { [Languages.bash]: { language: 'bash', version: '5.2.0', fileName: 'script.sh' }, @@ -259,6 +260,9 @@ export function resolveRuntimeSessionMode( ); } +const sandboxBackend = resolveSandboxBackend(process.env.CODEAPI_SANDBOX_BACKEND); +const runtimeSessionMode = resolveRuntimeSessionMode(process.env.CODEAPI_RUNTIME_SESSION_MODE); + export const env = { PORT: process.env.SERVICE_PORT ?? 3112, LOCAL_MODE: process.env.LOCAL_MODE === 'true', @@ -344,7 +348,7 @@ export const env = { * (current Kubernetes/libkrun sandbox-runner). * - `lambda-microvm`: AWS Lambda MicroVM backend. */ - SANDBOX_BACKEND: resolveSandboxBackend(process.env.CODEAPI_SANDBOX_BACKEND), + SANDBOX_BACKEND: sandboxBackend, /** * Runtime session affinity for stateful sandbox backends. * - `stateless` (default): no runtime sessions; `runtime_session_hint` ignored. @@ -353,7 +357,17 @@ export const env = { * - `strict`: same serialized session semantics, and a session hint is * required instead of degrading requests without one to stateless. */ - RUNTIME_SESSION_MODE: resolveRuntimeSessionMode(process.env.CODEAPI_RUNTIME_SESSION_MODE), + RUNTIME_SESSION_MODE: runtimeSessionMode, + /** + * Deployment identity used by trusted callers to route each agent to the + * intended execution stack. `default` is HTTP/stateless; `stateful` is + * Lambda MicroVM with session affinity. The startup policy rejects mixed + * tuples so an endpoint cannot claim one profile while running the other. + */ + EXECUTION_PROFILE: resolveExecutionProfile( + process.env.CODEAPI_EXECUTION_PROFILE, + runtimeSessionMode, + ), RUNTIME_SESSION_LOCK_WAIT_MS: configuredNumber( process.env.CODEAPI_RUNTIME_SESSION_LOCK_WAIT_MS, 15_000, diff --git a/service/src/enum/service.ts b/service/src/enum/service.ts index 84e8924..b936404 100644 --- a/service/src/enum/service.ts +++ b/service/src/enum/service.ts @@ -2,11 +2,6 @@ export enum Jobs { execute = 'execute', } -export enum Queues { - python = 'python-queue', - other = 'other-queue', -} - export enum Languages { bash = 'bash', js = 'js', diff --git a/service/src/execution-profile.test.ts b/service/src/execution-profile.test.ts new file mode 100644 index 0000000..32b1f6a --- /dev/null +++ b/service/src/execution-profile.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, test } from 'bun:test'; +import { + checkExecutionProfileExpectation, + queueNamesForExecutionProfile, + resolveExecutionProfile, +} from './execution-profile'; + +describe('execution profile resolution', () => { + test('preserves the HTTP/stateless default when unset', () => { + expect(resolveExecutionProfile(undefined, 'stateless')).toBe('default'); + }); + + test('recognizes an existing stateful deployment when unset', () => { + expect(resolveExecutionProfile(undefined, 'affinity')).toBe('stateful'); + expect(resolveExecutionProfile(undefined, 'strict')).toBe('stateful'); + }); + + test('lets API-only pods infer stateful from session mode without worker config', () => { + expect(resolveExecutionProfile(undefined, 'affinity')).toBe('stateful'); + }); + + test('accepts only the two public execution profiles', () => { + expect(resolveExecutionProfile('default', 'stateless')).toBe('default'); + expect(resolveExecutionProfile('stateful', 'affinity')).toBe('stateful'); + expect(() => resolveExecutionProfile('lambda', 'affinity')).toThrow( + 'CODEAPI_EXECUTION_PROFILE must be one of: default, stateful', + ); + expect(() => resolveExecutionProfile('', 'stateless')).toThrow( + 'CODEAPI_EXECUTION_PROFILE', + ); + }); +}); + +describe('execution profile queue isolation', () => { + test('keeps the legacy queue names for the default profile', () => { + expect(queueNamesForExecutionProfile('default')).toEqual({ + python: 'python-queue', + other: 'other-queue', + }); + }); + + test('uses separate queues for stateful workers', () => { + expect(queueNamesForExecutionProfile('stateful')).toEqual({ + python: 'stateful-python-queue', + other: 'stateful-other-queue', + }); + }); +}); + +describe('execution profile request assertion', () => { + test('allows callers that omit the assertion for backwards compatibility', () => { + expect(checkExecutionProfileExpectation(undefined, 'default')).toEqual({ ok: true }); + }); + + test('allows a matching expected profile', () => { + expect(checkExecutionProfileExpectation('stateful', 'stateful')).toEqual({ ok: true }); + }); + + test('returns a typed conflict before a mismatched request can be routed', () => { + expect(checkExecutionProfileExpectation('stateful', 'default')).toEqual({ + ok: false, + status: 409, + body: { + error: 'Expected the stateful execution profile, but reached default', + code: 'execution_profile_mismatch', + expected_profile: 'stateful', + actual_profile: 'default', + }, + }); + }); + + test('rejects invalid profile names instead of treating them as mismatches', () => { + expect(checkExecutionProfileExpectation('aws', 'default')).toMatchObject({ + ok: false, + status: 400, + body: { + code: 'invalid_execution_profile', + expected_profile: 'aws', + actual_profile: 'default', + }, + }); + }); +}); diff --git a/service/src/execution-profile.ts b/service/src/execution-profile.ts new file mode 100644 index 0000000..177acc2 --- /dev/null +++ b/service/src/execution-profile.ts @@ -0,0 +1,97 @@ +export const EXECUTION_PROFILES = ['default', 'stateful'] as const; + +export type ExecutionProfile = typeof EXECUTION_PROFILES[number]; + +export const EXPECTED_EXECUTION_PROFILE_HEADER = 'X-CodeAPI-Expected-Profile'; +export const EXECUTION_PROFILE_HEADER = 'X-CodeAPI-Execution-Profile'; + +export interface ExecutionProfileQueueNames { + python: string; + other: string; +} + +export function resolveExecutionProfile( + raw: string | undefined, + runtimeSessionMode: 'stateless' | 'affinity' | 'strict', +): ExecutionProfile { + if (raw != null) { + if (EXECUTION_PROFILES.includes(raw as ExecutionProfile)) { + return raw as ExecutionProfile; + } + throw new Error( + `CODEAPI_EXECUTION_PROFILE must be one of: ${EXECUTION_PROFILES.join(', ')}`, + ); + } + + /* Preserve the two supported pre-profile deployments during rollout. The + * common stateless stack remains `default`; a stateful API-only pod can + * infer `stateful` from its session mode even though worker-only backend + * credentials/config are intentionally absent. Worker startup separately + * verifies that this profile is backed by Lambda. */ + return runtimeSessionMode !== 'stateless' + ? 'stateful' + : 'default'; +} + +export function queueNamesForExecutionProfile( + profile: ExecutionProfile, +): ExecutionProfileQueueNames { + if (profile === 'stateful') { + return { + python: 'stateful-python-queue', + other: 'stateful-other-queue', + }; + } + return { + python: 'python-queue', + other: 'other-queue', + }; +} + +export type ExecutionProfileExpectation = + | { ok: true } + | { + ok: false; + status: 400 | 409; + body: { + error: string; + code: 'invalid_execution_profile' | 'execution_profile_mismatch'; + expected_profile?: string; + actual_profile: ExecutionProfile; + }; + }; + +export function checkExecutionProfileExpectation( + rawExpectedProfile: string | undefined, + actualProfile: ExecutionProfile, +): ExecutionProfileExpectation { + if (rawExpectedProfile == null) return { ok: true }; + + if (!EXECUTION_PROFILES.includes(rawExpectedProfile as ExecutionProfile)) { + return { + ok: false, + status: 400, + body: { + error: `Invalid execution profile: ${rawExpectedProfile}`, + code: 'invalid_execution_profile', + expected_profile: rawExpectedProfile, + actual_profile: actualProfile, + }, + }; + } + + if (rawExpectedProfile !== actualProfile) { + return { + ok: false, + status: 409, + body: { + error: `Expected the ${rawExpectedProfile} execution profile, but reached ${actualProfile}`, + code: 'execution_profile_mismatch', + expected_profile: rawExpectedProfile, + actual_profile: actualProfile, + }, + }; + } + + return { ok: true }; +} diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 59a18a8..e27846e 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -3,7 +3,12 @@ import type { Express } from 'express'; import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from './queue'; import { validateStartupAuthConfig } from './auth/startup'; import { env } from './config'; -import { validateApiHardenedConfig, validateSandboxBackendPolicy, validateWorkerHardenedConfig } from './secure-startup'; +import { + validateApiHardenedConfig, + validateExecutionProfilePolicy, + validateSandboxBackendPolicy, + validateWorkerHardenedConfig, +} from './secure-startup'; import logger from './logger'; import { shutdownTelemetry } from './telemetry'; @@ -75,6 +80,7 @@ function setupQueueListeners(queue: Queue, name: string): void { export async function startupApiOnly(): Promise { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); + validateExecutionProfilePolicy({ requireBackendMatch: false }); /* 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 @@ -97,6 +103,7 @@ export async function startupApiOnly(): Promise { export async function startupWorkerOnly(): Promise { logger.info('Starting Worker service...'); validateWorkerHardenedConfig(); + validateExecutionProfilePolicy(); validateSandboxBackendPolicy(); // Dynamically import workers to start them @@ -131,6 +138,7 @@ async function gracefulStartup(): Promise { logger.info('Starting up service (combined API + Workers)...'); validateApiHardenedConfig(); validateWorkerHardenedConfig(); + validateExecutionProfilePolicy(); validateSandboxBackendPolicy(); await validateLifecycleAuthConfig(); diff --git a/service/src/local-api.ts b/service/src/local-api.ts index b9b280d..7428346 100644 --- a/service/src/local-api.ts +++ b/service/src/local-api.ts @@ -11,6 +11,7 @@ import express, { json, Router } from 'express'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; import { requestErrorLogger, requestNotFoundLogger } from './middleware/request-error-logger'; +import { executionProfileMiddleware } from './middleware/execution-profile'; import { localAuth } from './auth/local'; import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from './queue'; import { setStartupComplete } from './lifecycle'; @@ -19,10 +20,12 @@ import './workers'; import { env } from './config'; import logger from './logger'; import { shutdownTelemetry, traceHttpRequest } from './telemetry'; +import { validateExecutionProfilePolicy } from './secure-startup'; const app = express(); app.disable('x-powered-by'); app.set('trust proxy', 1); +app.use(executionProfileMiddleware); let localShuttingDown = false; const v1 = Router(); @@ -52,6 +55,7 @@ app.use(requestErrorLogger); async function localStartup(): Promise { logger.info('Starting local development server...'); logger.info('⚠️ LOCAL MODE - No authentication required'); + validateExecutionProfilePolicy(); try { // Set a local user ID for session management diff --git a/service/src/metrics.ts b/service/src/metrics.ts index ba052dc..a2d9865 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -1,8 +1,27 @@ import client, { register, Counter, Histogram, Gauge } from 'prom-client'; import { normalizeMetricPath } from './httpPathNormalize'; +import { env } from './config'; client.collectDefaultMetrics({ register }); +export const executionProfileInfo = new Gauge({ + name: 'codeapi_execution_profile_info', + help: 'Static identity of this Code API execution deployment', + labelNames: ['profile', 'sandbox_backend', 'runtime_session_mode'] as const, +}); + +executionProfileInfo.set({ + profile: env.EXECUTION_PROFILE, + sandbox_backend: env.SANDBOX_BACKEND, + runtime_session_mode: env.RUNTIME_SESSION_MODE, +}, 1); + +export const executionProfileRequestRejections = new Counter({ + name: 'codeapi_execution_profile_request_rejections_total', + help: 'Requests rejected because the expected execution profile was invalid or mismatched', + labelNames: ['reason'] as const, +}); + // -- HTTP metrics (shared across Express and Bun servers) -- export const httpRequestsTotal = new Counter({ name: 'codeapi_http_requests_total', diff --git a/service/src/middleware/execution-profile.test.ts b/service/src/middleware/execution-profile.test.ts new file mode 100644 index 0000000..43368ef --- /dev/null +++ b/service/src/middleware/execution-profile.test.ts @@ -0,0 +1,78 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import type { NextFunction, Request, Response } from 'express'; +import { env } from '../config'; +import { executionProfileMiddleware } from './execution-profile'; + +const savedProfile = env.EXECUTION_PROFILE; + +afterEach(() => { + env.EXECUTION_PROFILE = savedProfile; +}); + +function invoke(expectedProfile?: string): { + headers: Record; + status?: number; + body?: unknown; + nextCalled: boolean; +} { + const result: { + headers: Record; + status?: number; + body?: unknown; + nextCalled: boolean; + } = { headers: {}, nextCalled: false }; + const req = { + get: () => expectedProfile, + } as unknown as Request; + const res = { + setHeader: (name: string, value: string) => { + result.headers[name] = value; + }, + status: (status: number) => { + result.status = status; + return res; + }, + json: (body: unknown) => { + result.body = body; + return res; + }, + } as unknown as Response; + const next = (() => { + result.nextCalled = true; + }) as NextFunction; + + executionProfileMiddleware(req, res, next); + return result; +} + +describe('execution profile middleware', () => { + test('advertises the actual profile and allows matching requests', () => { + env.EXECUTION_PROFILE = 'stateful'; + expect(invoke('stateful')).toEqual({ + headers: { 'X-CodeAPI-Execution-Profile': 'stateful' }, + nextCalled: true, + }); + }); + + test('rejects a mismatched endpoint before routing', () => { + env.EXECUTION_PROFILE = 'default'; + expect(invoke('stateful')).toMatchObject({ + headers: { 'X-CodeAPI-Execution-Profile': 'default' }, + status: 409, + body: { + code: 'execution_profile_mismatch', + expected_profile: 'stateful', + actual_profile: 'default', + }, + nextCalled: false, + }); + }); + + test('keeps older callers working when they omit the assertion', () => { + env.EXECUTION_PROFILE = 'default'; + expect(invoke()).toEqual({ + headers: { 'X-CodeAPI-Execution-Profile': 'default' }, + nextCalled: true, + }); + }); +}); diff --git a/service/src/middleware/execution-profile.ts b/service/src/middleware/execution-profile.ts new file mode 100644 index 0000000..2e3bdae --- /dev/null +++ b/service/src/middleware/execution-profile.ts @@ -0,0 +1,37 @@ +import type { NextFunction, Request, Response } from 'express'; +import { env } from '../config'; +import { + checkExecutionProfileExpectation, + EXECUTION_PROFILE_HEADER, + EXPECTED_EXECUTION_PROFILE_HEADER, +} from '../execution-profile'; +import { executionProfileRequestRejections } from '../metrics'; + +/** + * Advertise this deployment's profile and fail closed when a trusted caller + * reaches the wrong endpoint. Apply before routing so no file, programmatic, + * or ordinary execution request can enqueue work on a mismatched stack. + */ +export function executionProfileMiddleware( + req: Request, + res: Response, + next: NextFunction, +): void { + res.setHeader(EXECUTION_PROFILE_HEADER, env.EXECUTION_PROFILE); + + const expectation = checkExecutionProfileExpectation( + req.get(EXPECTED_EXECUTION_PROFILE_HEADER), + env.EXECUTION_PROFILE, + ); + if (expectation.ok) { + next(); + return; + } + + executionProfileRequestRejections.inc({ + reason: expectation.body.code === 'execution_profile_mismatch' + ? 'mismatch' + : 'invalid', + }); + res.status(expectation.status).json(expectation.body); +} diff --git a/service/src/queue.ts b/service/src/queue.ts index c6f3226..3410968 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -5,8 +5,9 @@ import { setMaxListeners } from 'events'; import type { CommonRedisOptions } from 'ioredis'; import type * as tls from 'tls'; import type * as t from './types'; -import { Jobs, Queues } from './enum'; +import { Jobs } from './enum'; import { env } from './config'; +import { queueNamesForExecutionProfile } from './execution-profile'; import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; @@ -54,16 +55,19 @@ const connection = new IORedis({ // Global queues - no INSTANCE_ID prefix // This enables horizontal scaling where any worker can process any job -const pyQueue = new Queue(Queues.python, { connection }); -const otherQueue = new Queue(Queues.other, { connection }); +// while the execution-profile prefix prevents HTTP and Lambda workers from +// consuming each other's jobs when they share Redis. +const queueNames = queueNamesForExecutionProfile(env.EXECUTION_PROFILE); +const pyQueue = new Queue(queueNames.python, { connection }); +const otherQueue = new Queue(queueNames.other, { connection }); -const pyQueueEvents = new QueueEvents(Queues.python, { connection }); -const otherQueueEvents = new QueueEvents(Queues.other, { connection }); +const pyQueueEvents = new QueueEvents(queueNames.python, { connection }); +const otherQueueEvents = new QueueEvents(queueNames.other, { connection }); const queueMetricStates = ['waiting', 'active', 'delayed'] as const; const queueMetricSources = [ - { name: Queues.python, queue: pyQueue }, - { name: Queues.other, queue: otherQueue }, + { name: queueNames.python, queue: pyQueue }, + { name: queueNames.other, queue: otherQueue }, ] as const; const QUEUE_METRICS_TIMEOUT_MS = 1000; @@ -109,4 +113,4 @@ registerBullmqQueueMetricsCollector(async () => { * BullMQ coordination objects. */ setMaxListeners(0, pyQueue, otherQueue, pyQueueEvents, otherQueueEvents); -export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection }; +export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, queueNames, connection }; diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 720809b..7e94960 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -3,6 +3,7 @@ import { env } from './config'; import { validateApiHardenedConfig, validateEgressGatewayHardenedConfig, + validateExecutionProfilePolicy, validateSandboxBackendPolicy, validateWorkerHardenedConfig, } from './secure-startup'; @@ -10,6 +11,7 @@ import { const savedEnv = { ...process.env }; const saved = { hardened: env.HARDENED_SANDBOX_MODE, + executionProfile: env.EXECUTION_PROFILE, sandboxBackend: env.SANDBOX_BACKEND, ptcMode: env.PTC_MODE, runtimeSessionMode: env.RUNTIME_SESSION_MODE, @@ -46,6 +48,7 @@ function restore(): void { } Object.assign(process.env, savedEnv); env.HARDENED_SANDBOX_MODE = saved.hardened; + env.EXECUTION_PROFILE = saved.executionProfile; env.SANDBOX_BACKEND = saved.sandboxBackend; env.PTC_MODE = saved.ptcMode; env.RUNTIME_SESSION_MODE = saved.runtimeSessionMode; @@ -78,6 +81,61 @@ function restore(): void { afterEach(restore); +describe('execution profile policy', () => { + test('accepts the AWS-free default profile', () => { + env.EXECUTION_PROFILE = 'default'; + env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'stateless'; + expect(() => validateExecutionProfilePolicy()).not.toThrow(); + }); + + test('accepts affinity and strict stateful profiles', () => { + env.EXECUTION_PROFILE = 'stateful'; + env.SANDBOX_BACKEND = 'lambda-microvm'; + env.RUNTIME_SESSION_MODE = 'affinity'; + expect(() => validateExecutionProfilePolicy()).not.toThrow(); + env.RUNTIME_SESSION_MODE = 'strict'; + expect(() => validateExecutionProfilePolicy()).not.toThrow(); + }); + + test('does not require worker-only backend config on API-only pods', () => { + env.EXECUTION_PROFILE = 'stateful'; + env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'affinity'; + expect(() => validateExecutionProfilePolicy({ requireBackendMatch: false })).not.toThrow(); + }); + + test('rejects a default profile backed by AWS or stateful sessions', () => { + env.EXECUTION_PROFILE = 'default'; + env.SANDBOX_BACKEND = 'lambda-microvm'; + env.RUNTIME_SESSION_MODE = 'stateless'; + expect(() => validateExecutionProfilePolicy()).toThrow( + 'CODEAPI_EXECUTION_PROFILE=default requires', + ); + + env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'affinity'; + expect(() => validateExecutionProfilePolicy()).toThrow( + 'CODEAPI_EXECUTION_PROFILE=default requires', + ); + }); + + test('rejects a stateful profile without Lambda affinity', () => { + env.EXECUTION_PROFILE = 'stateful'; + env.SANDBOX_BACKEND = 'http'; + env.RUNTIME_SESSION_MODE = 'affinity'; + expect(() => validateExecutionProfilePolicy()).toThrow( + 'CODEAPI_EXECUTION_PROFILE=stateful requires', + ); + + env.SANDBOX_BACKEND = 'lambda-microvm'; + env.RUNTIME_SESSION_MODE = 'stateless'; + expect(() => validateExecutionProfilePolicy()).toThrow( + 'CODEAPI_EXECUTION_PROFILE=stateful requires', + ); + }); +}); + describe('hardened CodeAPI startup config', () => { test('rejects grant secrets in API and worker processes', () => { env.HARDENED_SANDBOX_MODE = true; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 8d4729d..df9409e 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -63,6 +63,41 @@ export function validateWorkerHardenedConfig(): void { requireValue('CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY', env.EXECUTION_MANIFEST_PRIVATE_KEY); } +/** + * Make the endpoint identity trustworthy. Callers route by execution profile, + * so accepting a contradictory backend/session tuple would silently send work + * to the wrong infrastructure and could lose workspace continuity. + */ +export function validateExecutionProfilePolicy(options: { + requireBackendMatch?: boolean; +} = {}): void { + const requireBackendMatch = options.requireBackendMatch ?? true; + if (env.EXECUTION_PROFILE === 'default') { + if ( + env.RUNTIME_SESSION_MODE !== 'stateless' + || (requireBackendMatch && env.SANDBOX_BACKEND !== 'http') + ) { + throw new SecureStartupConfigError( + 'CODEAPI_EXECUTION_PROFILE=default requires ' + + (requireBackendMatch ? 'CODEAPI_SANDBOX_BACKEND=http and ' : '') + + 'CODEAPI_RUNTIME_SESSION_MODE=stateless', + ); + } + return; + } + + if ( + env.RUNTIME_SESSION_MODE === 'stateless' + || (requireBackendMatch && env.SANDBOX_BACKEND !== 'lambda-microvm') + ) { + throw new SecureStartupConfigError( + 'CODEAPI_EXECUTION_PROFILE=stateful requires ' + + (requireBackendMatch ? 'CODEAPI_SANDBOX_BACKEND=lambda-microvm and ' : '') + + 'CODEAPI_RUNTIME_SESSION_MODE=affinity or strict', + ); + } +} + /** * Backend-selection policy. Unlike the hardened-mode validators, this runs * unconditionally: a misconfigured backend must never half-start. diff --git a/service/src/service-api.ts b/service/src/service-api.ts index c664f14..ec15b1a 100644 --- a/service/src/service-api.ts +++ b/service/src/service-api.ts @@ -2,6 +2,7 @@ import express, { json, Router } from 'express'; import { startServer, gracefulShutdown } from './lifecycle'; import { apiKeyAuth } from './middleware/auth'; import { requestErrorLogger, requestNotFoundLogger } from './middleware/request-error-logger'; +import { executionProfileMiddleware } from './middleware/execution-profile'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; import { connection } from './queue'; @@ -11,6 +12,7 @@ import logger from './logger'; const app = express(); app.disable('x-powered-by'); app.set('trust proxy', 1); +app.use(executionProfileMiddleware); const v1 = Router(); diff --git a/service/src/service/router.ts b/service/src/service/router.ts index d88dfdc..4c844fb 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -10,7 +10,7 @@ import { sessionAuth } from '../middleware/auth'; import { executionLimiter, uploadLimiter, downloadLimiter, fetchLimiter } from '../middleware/limits'; import { internalServiceHeaders } from '../internal-service-auth'; import { resolveSessionKey, resolveOutputBucketSessionKey, SessionKeyResolutionError, parseUploadSessionKeyInput, type SessionKeyInput } from '../session-key'; -import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from '../queue'; +import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, queueNames, connection } from '../queue'; import { sleep, getAxiosErrorDetails, publicExecutionFailure } from '../utils'; import { env, jobCompletionWaitTimeoutMs, planLimits, resolveLanguage } from '../config'; import { createPayload } from '../payload'; @@ -226,12 +226,13 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) const queue = language === Languages.py ? pyQueue : otherQueue; const queueEvents = language === Languages.py ? pyQueueEvents : otherQueueEvents; - const queueName = language === Languages.py ? 'python' : 'other'; + const queueName = language === Languages.py ? queueNames.python : queueNames.other; const job = await withSpan('codeapi.job.enqueue', { 'messaging.system': 'bullmq', 'messaging.destination.name': queueName, 'codeapi.language': language, + 'codeapi.execution_profile': env.EXECUTION_PROFILE, }, () => { const traceCarrier = captureTraceCarrier(); return queue.add(Jobs.execute, { @@ -279,6 +280,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) 'messaging.system': 'bullmq', 'messaging.destination.name': queueName, 'codeapi.language': language, + 'codeapi.execution_profile': env.EXECUTION_PROFILE, }, () => job.waitUntilFinished(queueEvents, JOB_COMPLETION_WAIT_TIMEOUT_MS), 'CONSUMER'); if (!isSyntheticRequest) { diff --git a/service/src/workers.ts b/service/src/workers.ts index e9c3b2b..206542b 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -3,8 +3,7 @@ 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 { Queues } from './enum'; -import { connection } from './queue'; +import { connection, queueNames } from './queue'; import { env, jobDeadlineAtMs } from './config'; import { summarizeSandboxResponse, summarizeText } from './execution-log'; import { createGatewayEgressGrant, restoreGatewaySandboxResult, revokeGatewayEgressGrant } from './egress-gateway-client'; @@ -238,7 +237,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { // Global workers - no INSTANCE_ID prefix // This enables horizontal scaling where any worker can process any job from the shared queue // Each worker respects its own concurrency limit based on its co-located sandbox capacity -export const pyWorker = new Worker(Queues.python, processJob, { +export const pyWorker = new Worker(queueNames.python, processJob, { connection, concurrency: env.PYTHON_CONCURRENCY, limiter: { @@ -247,7 +246,7 @@ export const pyWorker = new Worker(Queues.python, processJob, { }, }); -export const otherWorker = new Worker(Queues.other, processJob, { +export const otherWorker = new Worker(queueNames.other, processJob, { connection, concurrency: env.OTHER_CONCURRENCY, limiter: {