diff --git a/.env.example b/.env.example index 54a2971a..3bd113e5 100644 --- a/.env.example +++ b/.env.example @@ -24,6 +24,13 @@ SANDBOX_RUN_CPU_TIME=10000 SANDBOX_RUN_TIMEOUT=15000 SANDBOX_OUTPUT_MAX_SIZE=65536 +# Remote stateful code bridge (Code API deployment) +# CODEAPI_SANDBOX_BACKEND=remote-bridge +# CODEAPI_EXECUTION_PROFILE=stateful +# CODEAPI_RUNTIME_SESSION_MODE=affinity +# CODEAPI_BRIDGE_WORKER_ID=my-vm +# CODEAPI_BRIDGE_TOKEN=replace-with-a-strong-random-secret + # Service Configuration PYTHON_CONCURRENCY=5 OTHER_CONCURRENCY=15 diff --git a/.gitignore b/.gitignore index db2b8a28..a1a0c6ed 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ data/ node_modules +packages/*/dist/ .env .git .npmrc diff --git a/README.md b/README.md index 6612cb5f..716941d9 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Code Interpreter (internally `codeapi`, the prefix used by its env vars, images, - **Package Delivery** - Bakes Python, Node, and Bun into the default microVM block-root image; a package-init PVC mode remains available for direct NsJail development +- **Remote Code Bridge** - Lets an operator-owned VM connect outbound and serve + as a fenced, stateful sandbox through the `@librechat/code` worker ## Architecture @@ -65,6 +67,18 @@ Two modes are supported: - **NsJail mode** (`kvmEnabled: false`): Direct NsJail sandboxing with Linux namespaces and cgroups - **MicroVM mode** (`kvmEnabled: true`): libkrun microVM with its own kernel, NsJail runs inside the guest +## Remote stateful environments + +The `remote-bridge` backend keeps the Code API as the policy and queue boundary +while moving execution to a sandbox on an operator-selected VM. The worker only +makes outbound authenticated requests, so the VM does not need a public ingress +port. Assignments carry a deadline, a single-active-worker lock, a monotonically +increasing generation, and a one-time lease token to fence stale workers. + +See [Remote Code Bridge](docs/remote-bridge/README.md) for deployment and threat +model details. The worker protocol and CLI live in the provider-neutral +[`@librechat/code`](packages/code/README.md) package. + ## Security disclaimer This service exists to run arbitrary, untrusted code — treat every diff --git a/api/Dockerfile b/api/Dockerfile index f8d6713c..3526e98e 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -29,8 +29,10 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \ RUN make -j$(nproc) COPY api/src/spec-guard.c /tmp/spec-guard.c +COPY docker/rootfs-setup.c /tmp/rootfs-setup.c RUN gcc -O2 -static -o /usr/local/bin/spec-guard /tmp/spec-guard.c \ - && chmod 0111 /usr/local/bin/spec-guard + && gcc -O2 -static -o /usr/local/bin/sandbox-rootfs-setup /tmp/rootfs-setup.c \ + && chmod 0111 /usr/local/bin/spec-guard /usr/local/bin/sandbox-rootfs-setup # ============================================================================ # Stage 1b: Build language runtime packages (only consumed by sandbox-runner-baked) @@ -212,6 +214,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ && dnf clean all COPY --from=launcher-builder /launcher/target/release/sandbox-launcher /usr/local/bin/launcher +COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-setup COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh COPY docker/start-direct-sandbox.sh /usr/local/bin/start-direct-sandbox.sh diff --git a/api/src/entrypoint.sh b/api/src/entrypoint.sh index b532a67c..4fce5ba4 100755 --- a/api/src/entrypoint.sh +++ b/api/src/entrypoint.sh @@ -171,28 +171,33 @@ fi chmod 777 "$SMOKE_DIR" fi SMOKE_LOG=$(mktemp) +SMOKE_STDERR=$(mktemp) NSJAIL_CGROUP_ARGS=() if [ "$SANDBOX_USE_CGROUPV2" = "true" ]; then NSJAIL_CGROUP_ARGS=(--use_cgroupv2) fi -if timeout 10 /usr/sbin/nsjail --config "${NSJAIL_CONFIG:-/sandbox_api/config/sandbox.cfg}" \ +if timeout 10 "${NSJAIL_PATH:-/usr/sbin/nsjail}" --config "${NSJAIL_CONFIG:-/sandbox_api/config/sandbox.cfg}" \ "${NSJAIL_CGROUP_ARGS[@]}" --log "$SMOKE_LOG" \ --user "65534:${SMOKE_OUTSIDE_UID}:1" --group "65534:${SMOKE_OUTSIDE_GID}:1" \ -s /usr/bin:/bin -s /usr/lib:/lib -s /usr/lib64:/lib64 \ -B "$SMOKE_DIR:/mnt/data" \ - -- /bin/sh -c 'printf "%s\n" sandbox_ok > /mnt/data/smoke.txt && test "$(cat /mnt/data/smoke.txt)" = sandbox_ok' > /dev/null 2>&1; then + -- /bin/sh -c 'printf "%s\n" sandbox_ok > /mnt/data/smoke.txt && test "$(cat /mnt/data/smoke.txt)" = sandbox_ok' > /dev/null 2>"$SMOKE_STDERR"; then echo "NsJail smoke test passed" else echo "FATAL: NsJail smoke test failed — sandbox cannot start" echo "NsJail log output:" cat "$SMOKE_LOG" 2>/dev/null || true + echo "NsJail stderr:" + cat "$SMOKE_STDERR" 2>/dev/null || true rm -f "$SMOKE_LOG" + rm -f "$SMOKE_STDERR" rm -rf "$SMOKE_DIR" exit 1 fi rm -f "$SMOKE_LOG" +rm -f "$SMOKE_STDERR" rm -rf "$SMOKE_DIR" echo "Starting sandbox API server..." diff --git a/docker/Dockerfile.worker-sandbox b/docker/Dockerfile.worker-sandbox index cd3edd3d..c18eab82 100644 --- a/docker/Dockerfile.worker-sandbox +++ b/docker/Dockerfile.worker-sandbox @@ -42,8 +42,10 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \ RUN make -j$(nproc) COPY api/src/spec-guard.c /tmp/spec-guard.c +COPY docker/rootfs-setup.c /tmp/rootfs-setup.c RUN gcc -O2 -static -o /usr/local/bin/spec-guard /tmp/spec-guard.c \ - && chmod 0111 /usr/local/bin/spec-guard + && gcc -O2 -static -o /usr/local/bin/sandbox-rootfs-setup /tmp/rootfs-setup.c \ + && chmod 0111 /usr/local/bin/spec-guard /usr/local/bin/sandbox-rootfs-setup # ============================================================================ # Stage 1b: Build language runtime packages for the baked KVM root disk @@ -83,6 +85,7 @@ WORKDIR /app COPY service/package.json service/bun.lock ./ RUN bun install --frozen-lockfile COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ RUN bun build ./src/worker-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' @@ -231,6 +234,7 @@ ENV PATH="/root/.bun/bin:${PATH}" # --- Launcher (runs on host, boots microVM) --- COPY --from=launcher-builder /launcher/target/release/sandbox-launcher /usr/local/bin/launcher +COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-setup # --- Launcher entrypoint (DNS resolution + socat relay before VM boot) --- COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh diff --git a/docker/rootfs-setup.c b/docker/rootfs-setup.c new file mode 100644 index 00000000..b0b817cf --- /dev/null +++ b/docker/rootfs-setup.c @@ -0,0 +1,92 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +static int bind_mount(const char *source, const char *target, int read_only) { + if (mount(source, target, NULL, MS_BIND | MS_REC, NULL) != 0) { + fprintf(stderr, "bind %s -> %s failed: %s\n", source, target, strerror(errno)); + return -1; + } + + if (read_only && + mount(NULL, target, NULL, MS_BIND | MS_REMOUNT | MS_RDONLY, NULL) != 0) { + fprintf(stderr, "read-only remount of %s failed: %s\n", target, strerror(errno)); + return -1; + } + + return 0; +} + +static int bind_rootfs_path(const char *rootfs, const char *path) { + char source[PATH_MAX]; + int written = snprintf(source, sizeof(source), "%s%s", rootfs, path); + if (written < 0 || (size_t)written >= sizeof(source)) { + fprintf(stderr, "rootfs path is too long: %s%s\n", rootfs, path); + return -1; + } + + return bind_mount(source, path, 1); +} + +int main(int argc, char **argv) { + if (argc < 2) { + fprintf(stderr, "usage: sandbox-rootfs-setup ROOTFS [COMMAND ...]\n"); + return 2; + } + + const char *rootfs = argv[1]; + if (rootfs[0] != '/') { + fprintf(stderr, "rootfs must be an absolute path\n"); + return 2; + } + + if (mount(NULL, "/", NULL, MS_REC | MS_PRIVATE, NULL) != 0) { + fprintf(stderr, "making the mount namespace private failed: %s\n", strerror(errno)); + return 1; + } + + if ((mkdir("/sandbox_api", 0755) != 0 && errno != EEXIST) || + (mkdir("/pkgs", 0755) != 0 && errno != EEXIST)) { + fprintf(stderr, "creating rootfs mount targets failed: %s\n", strerror(errno)); + return 1; + } + + /* + * Keep this process statically linked: the final /usr mount replaces + * the Fedora launcher's dynamic userspace with the Debian sandbox rootfs. + * A shell cannot safely perform this sequence because its next command may + * try to load a host binary against guest libraries (or vice versa). + */ + const char *paths[] = {"/sandbox_api", "/pkgs"}; + for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); i++) { + if (bind_rootfs_path(rootfs, paths[i]) != 0) { + return 1; + } + } + + if (access("/host-packages", F_OK) == 0 && + bind_mount("/host-packages", "/pkgs", 0) != 0) { + fprintf(stderr, "warning: sandbox will run without host packages\n"); + } + + /* Bind all guest userspace last, then immediately enter it. */ + if (bind_rootfs_path(rootfs, "/usr") != 0) { + return 1; + } + + setenv("PATH", "/root/.bun/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", 1); + setenv("LD_LIBRARY_PATH", "/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu", 1); + + setenv("NSJAIL_PATH", "/usr/sbin/nsjail", 1); + + char *default_argv[] = {"/sandbox_api/entrypoint.sh", NULL}; + char **command_argv = argc > 2 ? &argv[2] : default_argv; + execv(command_argv[0], command_argv); + fprintf(stderr, "starting sandbox entrypoint failed: %s\n", strerror(errno)); + return 1; +} diff --git a/docker/start-direct-sandbox.sh b/docker/start-direct-sandbox.sh index a171a6df..bf7c9805 100644 --- a/docker/start-direct-sandbox.sh +++ b/docker/start-direct-sandbox.sh @@ -44,35 +44,4 @@ else fi export SANDBOX_ROOTFS="$ROOTFS" - -exec unshare --mount bash -c ' - ROOTFS="${SANDBOX_ROOTFS:-/sandbox-rootfs}" - - mount -o bind,ro "$ROOTFS/usr/sbin" /usr/sbin || { echo "FATAL: cannot bind /usr/sbin"; exit 1; } - mount -o bind,ro "$ROOTFS/usr/lib" /usr/lib || { echo "FATAL: cannot bind /usr/lib"; exit 1; } - - if [ -d "$ROOTFS/usr/lib64" ] && ! [ -L "$ROOTFS/usr/lib64" ]; then - mount -o bind,ro "$ROOTFS/usr/lib64" /usr/lib64 2>/dev/null || \ - echo "[sandbox] WARNING: could not bind /usr/lib64 - sandboxed binaries may fail to exec" - fi - - mount -o bind,ro "$ROOTFS/usr/local" /usr/local || { echo "FATAL: cannot bind /usr/local"; exit 1; } - mount -o bind,ro "$ROOTFS/sandbox_api" /sandbox_api || { echo "FATAL: cannot bind /sandbox_api"; exit 1; } - mount -o bind,ro "$ROOTFS/pkgs" /pkgs || { echo "FATAL: cannot bind /pkgs"; exit 1; } - - if [ -d /host-packages ]; then - mount --bind /host-packages /pkgs 2>/dev/null || \ - echo "WARNING: could not bind /host-packages - sandbox will run without packages" - fi - - mount -o bind,ro "$ROOTFS/usr/bin" /usr/bin || { echo "FATAL: cannot bind /usr/bin"; exit 1; } - - multiarch_libdir=$(find /usr/lib -maxdepth 1 -type d -name "*-linux-gnu" -print -quit) - if [ -n "$multiarch_libdir" ]; then - export LD_LIBRARY_PATH="$multiarch_libdir${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" - fi - - export PATH="/root/.bun/bin:$PATH" - - exec /sandbox_api/entrypoint.sh -' +exec unshare --mount /sandbox-rootfs-setup "$ROOTFS" diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md new file mode 100644 index 00000000..a8d50e2a --- /dev/null +++ b/docs/remote-bridge/README.md @@ -0,0 +1,94 @@ +# Remote Code Bridge + +Remote Code Bridge makes an operator-owned VM a stateful Code API execution +environment without exposing that VM to inbound internet traffic. + +```text +LibreChat -> Code API -> Redis assignment + ^ | + | outbound v + @librechat/code -> local sandbox +``` + +Code API remains the public authentication, policy, manifest, timeout, and +result-normalization boundary. The bridge worker has a separate operator +credential and never accepts end-user bearer tokens directly. + +## Code API configuration + +Run this as an isolated stateful Code API deployment: + +```dotenv +CODEAPI_SANDBOX_BACKEND=remote-bridge +CODEAPI_EXECUTION_PROFILE=stateful +CODEAPI_RUNTIME_SESSION_MODE=affinity +CODEAPI_BRIDGE_WORKER_ID=my-vm +CODEAPI_BRIDGE_TOKEN= +``` + +Use `strict` instead of `affinity` if every request must include a runtime +session hint. In hardened mode, startup requires the bridge token to be at least +32 bytes. `PTC_MODE=blocking` is rejected; replay mode is required because a +remote execution cannot retain an open Code API process across tool callbacks. + +Start the CLI beside a sandbox using the same worker ID and secret; see +[`@librechat/code`](../../packages/code/README.md). +Stateful deployments must also set `LIBRECHAT_CODE_STATEFUL_WORKSPACE=true` +and route the CLI's `{runtimeSessionId}` endpoint template to an isolated, +persistent local runner per session. A single sandbox endpoint is stateless and +is rejected for runtime-session assignments. + +## LibreChat configuration + +Expose the Code API deployment as an environment under the Agents endpoint: + +```yaml +endpoints: + agents: + statefulCodeSessions: + environments: + - id: my-vm + name: My VM + type: attached + baseURL: https://code.example.com/v1 + default: true +``` + +Agents may select this environment with `code_environment_id: my-vm`. +LibreChat derives a stable per-conversation runtime session ID, so commands in +later turns reuse the same workspace. Attached environments deliberately skip +background prewarming: the single worker lease is reserved for explicit user +execution. + +## Lifecycle and fencing + +- Registration is ephemeral in Redis and must be refreshed by the worker. +- Code API permits one active assignment per configured worker. +- Each assignment has an absolute deadline, generation, and random lease token. +- Settlements with the wrong worker, generation, token, or expired deadline are + rejected. +- Assignments are queued for the exact registered worker incarnation, so an + outstanding poll from a replaced process cannot consume replacement work. +- Assignment records and the worker lock live through the full configured job + deadline plus cleanup grace. +- Ambiguous settlement delivery is retried through the assignment deadline. If + a stateful settlement remains ambiguous, the CLI exits and the affected local + session runner must be reset or discarded before restart. +- Request cancellation is polled by the worker and aborts the local sandbox + request. +- The sandbox receives the stable runtime session ID separately from the lease; + workspace state belongs to that session, not to a transient assignment. + +## Security boundaries + +The bridge removes inbound VM exposure; it does not replace sandbox isolation. +For internet-facing LibreChat deployments, use the hardened microVM/NsJail +stack, default-deny sandbox egress, signed execution manifests, least-privilege +host credentials, resource limits, and host/network monitoring. Bind the local +sandbox endpoint to loopback or a private container network. Rotate a leaked +bridge token immediately; the initial protocol intentionally uses a static +operator secret and supports one configured worker per Code API deployment. + +The next control-plane layer can add short-lived pairing credentials and a +multi-worker directory without changing the execution protocol or moving code +tools into the Agents SDK. diff --git a/launcher/Dockerfile b/launcher/Dockerfile index 0d252c28..1a077e12 100644 --- a/launcher/Dockerfile +++ b/launcher/Dockerfile @@ -33,8 +33,10 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \ RUN make -j$(nproc) COPY api/src/spec-guard.c /tmp/spec-guard.c +COPY docker/rootfs-setup.c /tmp/rootfs-setup.c RUN gcc -O2 -static -o /usr/local/bin/spec-guard /tmp/spec-guard.c \ - && chmod 0111 /usr/local/bin/spec-guard + && gcc -O2 -static -o /usr/local/bin/sandbox-rootfs-setup /tmp/rootfs-setup.c \ + && chmod 0111 /usr/local/bin/spec-guard /usr/local/bin/sandbox-rootfs-setup FROM oven/bun:1.3.14-debian AS sandbox-build @@ -126,6 +128,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ && dnf clean all COPY --from=launcher-builder /launcher/target/release/sandbox-launcher /usr/local/bin/launcher +COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-setup COPY --from=sandbox-build / /sandbox-rootfs/ @@ -136,6 +139,6 @@ RUN mkdir -p /host-packages COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh COPY docker/start-direct-sandbox.sh /usr/local/bin/start-direct-sandbox.sh COPY docker/sandbox-entrypoint.sh /usr/local/bin/entrypoint.sh -RUN chmod +x /usr/local/bin/launcher-entrypoint.sh /usr/local/bin/start-direct-sandbox.sh /usr/local/bin/entrypoint.sh +RUN chmod +x /sandbox-rootfs-setup /usr/local/bin/launcher-entrypoint.sh /usr/local/bin/start-direct-sandbox.sh /usr/local/bin/entrypoint.sh ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] diff --git a/packages/code/Dockerfile b/packages/code/Dockerfile new file mode 100644 index 00000000..21fb09bd --- /dev/null +++ b/packages/code/Dockerfile @@ -0,0 +1,15 @@ +FROM node:24-alpine AS build +WORKDIR /app +COPY package.json package-lock.json tsconfig.json ./ +RUN npm ci +COPY src ./src +RUN npm run build + +FROM node:24-alpine +ENV NODE_ENV=production +RUN addgroup -S librechat-code && adduser -S librechat-code -G librechat-code +WORKDIR /app +COPY --from=build /app/package.json ./package.json +COPY --from=build /app/dist ./dist +USER librechat-code +ENTRYPOINT ["node", "dist/cli.js"] diff --git a/packages/code/README.md b/packages/code/README.md new file mode 100644 index 00000000..edb8de2a --- /dev/null +++ b/packages/code/README.md @@ -0,0 +1,51 @@ +# `@librechat/code` + +Provider-neutral protocol and worker CLI for attaching a stateful, sandboxed +code environment to LibreChat Code API. + +The CLI is a transport bridge, not a sandbox. Run it beside a Code Interpreter +sandbox (NsJail for trusted local development, or the hardened microVM stack for +untrusted internet traffic). It connects outbound to Code API, long-polls for +assignments, forwards them to the local sandbox, and returns fenced results. +The VM does not need an inbound public port. + +## Run + +```bash +npm install -g @librechat/code + +LIBRECHAT_CODE_URL=https://code.example.com/v1 \ +LIBRECHAT_CODE_WORKER_TOKEN='' \ +LIBRECHAT_CODE_WORKER_ID=my-vm \ +LIBRECHAT_CODE_SANDBOX_ENDPOINT=http://127.0.0.1:2000/api/v2 \ +librechat-code +``` + +Optional environment variables: + +- `LIBRECHAT_CODE_SANDBOX_PROFILE`: capability label; defaults to `nsjail`. +- `LIBRECHAT_CODE_RUNTIMES`: comma-separated capability labels. +- `LIBRECHAT_CODE_POLICY`: local policy description hashed into the worker's + registration; defaults to `default-deny`. +- `LIBRECHAT_CODE_STATEFUL_WORKSPACE`: defaults to `false`. Set it to `true` + only when the local sandbox supervisor provides a distinct persistent runner + for every runtime session. In that mode the endpoint must contain a + `{runtimeSessionId}` placeholder, for example + `http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2`. The worker URL- + encodes and substitutes the assigned session ID before execution. Hintless + assignments use an ephemeral `assignment-` session so affinity-mode + stateless work never reaches a literal placeholder route. + +A single built-in sandbox runner binds itself to one runtime session and must +not be advertised as stateful. Use the default stateless capability until a +session-routing supervisor is configured. + +The worker retries result settlement through the assignment deadline. If a +stateful result remains ambiguous, it exits with a quarantine error instead of +accepting another assignment. Reset or discard that session's local runner +before restarting the worker; its workspace may contain mutations that Code +API did not commit. + +Use a unique worker ID and secret per Code API deployment, expose only the +sandbox loopback endpoint to the CLI, and enforce VM/container egress policy +independently of the bridge transport. diff --git a/packages/code/package-lock.json b/packages/code/package-lock.json new file mode 100644 index 00000000..15ca1ad6 --- /dev/null +++ b/packages/code/package-lock.json @@ -0,0 +1,54 @@ +{ + "name": "@librechat/code", + "version": "0.1.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "@librechat/code", + "version": "0.1.0", + "license": "Apache-2.0", + "bin": { + "librechat-code": "dist/cli.js" + }, + "devDependencies": { + "@types/node": "^22.5.5", + "typescript": "^5.5.4" + }, + "engines": { + "node": ">=20" + } + }, + "node_modules/@types/node": { + "version": "22.20.1", + "resolved": "https://registry.npmjs.org/@types/node/-/node-22.20.1.tgz", + "integrity": "sha512-EANqOCF9QFyra+4pfxUcX9STKJpCLjMbObVzljIJomAWSnuSIEAvyzEU53GaajbXJEgdh0iEcPL+DGvpUd4k1Q==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/typescript": { + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/packages/code/package.json b/packages/code/package.json new file mode 100644 index 00000000..f195881e --- /dev/null +++ b/packages/code/package.json @@ -0,0 +1,42 @@ +{ + "name": "@librechat/code", + "version": "0.1.0", + "description": "LibreChat stateful code environment protocol and worker CLI", + "license": "Apache-2.0", + "type": "module", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + }, + "./protocol": { + "types": "./dist/protocol.d.ts", + "import": "./dist/protocol.js" + }, + "./worker": { + "types": "./dist/worker.d.ts", + "import": "./dist/worker.js" + } + }, + "bin": { + "librechat-code": "./dist/cli.js" + }, + "files": [ + "dist", + "!dist/*.test.*" + ], + "scripts": { + "build": "tsc -p tsconfig.json", + "test": "npm run build && node --test dist/*.test.js", + "prepack": "npm run build" + }, + "devDependencies": { + "@types/node": "^22.5.5", + "typescript": "^5.5.4" + }, + "engines": { + "node": ">=20" + } +} diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts new file mode 100644 index 00000000..411dda7f --- /dev/null +++ b/packages/code/src/cli.ts @@ -0,0 +1,55 @@ +#!/usr/bin/env node +import { createHash } from 'node:crypto'; +import { BridgeWorker } from './worker.js'; + +function required(name: string): string { + const value = process.env[name]?.trim(); + if (!value) throw new Error(`${name} is required`); + return value; +} + +function list(value: string | undefined): string[] { + return ( + value + ?.split(',') + .map((item) => item.trim()) + .filter(Boolean) ?? [] + ); +} + +const controller = new AbortController(); +process.once('SIGINT', () => controller.abort()); +process.once('SIGTERM', () => controller.abort()); + +const policy = process.env.LIBRECHAT_CODE_POLICY ?? 'default-deny'; +const statefulWorkspace = + process.env.LIBRECHAT_CODE_STATEFUL_WORKSPACE?.trim().toLowerCase() === 'true'; +const sandboxEndpoint = + process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? + 'http://127.0.0.1:2000/api/v2'; +if (statefulWorkspace && !sandboxEndpoint.includes('{runtimeSessionId}')) { + throw new Error( + 'LIBRECHAT_CODE_STATEFUL_WORKSPACE requires LIBRECHAT_CODE_SANDBOX_ENDPOINT to contain {runtimeSessionId}', + ); +} +const worker = new BridgeWorker({ + codeApiUrl: required('LIBRECHAT_CODE_URL'), + token: required('LIBRECHAT_CODE_WORKER_TOKEN'), + workerId: required('LIBRECHAT_CODE_WORKER_ID'), + sandboxEndpoint, + capabilities: { + statefulWorkspace, + sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? 'nsjail', + runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), + policyDigest: createHash('sha256').update(policy).digest('hex'), + }, + onError: (error) => { + const message = error instanceof Error ? error.message : 'unknown bridge error'; + process.stderr.write(`librechat-code: reconnecting after ${message}\n`); + }, +}); + +worker.run(controller.signal).catch((error: Error) => { + process.stderr.write(`librechat-code: ${error.message}\n`); + process.exitCode = 1; +}); diff --git a/packages/code/src/index.ts b/packages/code/src/index.ts new file mode 100644 index 00000000..c5eeaafc --- /dev/null +++ b/packages/code/src/index.ts @@ -0,0 +1,2 @@ +export * from './protocol.js'; +export * from './worker.js'; diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts new file mode 100644 index 00000000..9223b3ab --- /dev/null +++ b/packages/code/src/protocol.test.ts @@ -0,0 +1,10 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import { bridgeWorkerPath } from './protocol.js'; + +test('bridgeWorkerPath encodes worker-controlled path segments', () => { + assert.equal( + bridgeWorkerPath('vm/example worker'), + '/bridge/workers/vm%2Fexample%20worker', + ); +}); diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts new file mode 100644 index 00000000..a79d4ff4 --- /dev/null +++ b/packages/code/src/protocol.ts @@ -0,0 +1,92 @@ +export const BRIDGE_PROTOCOL_VERSION = 1 as const; + +export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; + +export interface BridgeWorkerCapabilities { + statefulWorkspace: boolean; + sandboxProfile: string; + runtimes: string[]; + policyDigest?: string; +} + +export interface BridgeWorkerRegistration { + protocolVersion: BridgeProtocolVersion; + workerId: string; + incarnationId: string; + capabilities: BridgeWorkerCapabilities; +} + +export interface BridgeWorkerRegistrationResponse { + protocolVersion: BridgeProtocolVersion; + workerId: string; + incarnationId: string; + registeredAt: string; + leaseTtlMs: number; +} + +export interface BridgeSandboxRequest { + body: TBody; + headers: Record; +} + +export interface BridgeAssignment { + protocolVersion: BridgeProtocolVersion; + assignmentId: string; + workerId: string; + incarnationId: string; + generation: number; + leaseToken: string; + expiresAt: string; + runtimeSessionId?: string; + request: BridgeSandboxRequest; +} + +export interface BridgeLeaseResponse { + protocolVersion: BridgeProtocolVersion; + assignment?: BridgeAssignment; +} + +export interface BridgeFulfilledSettlement { + protocolVersion: BridgeProtocolVersion; + generation: number; + leaseToken: string; + incarnationId: string; + status: 'fulfilled'; + result: TResult; +} + +export interface BridgeRejectedSettlement { + protocolVersion: BridgeProtocolVersion; + generation: number; + leaseToken: string; + incarnationId: string; + status: 'rejected'; + error: string; +} + +export type BridgeSettlement = + BridgeFulfilledSettlement | BridgeRejectedSettlement; + +export interface BridgeSettlementResponse { + protocolVersion: BridgeProtocolVersion; + accepted: true; +} + +export interface BridgeCancellationResponse { + protocolVersion: BridgeProtocolVersion; + cancelled: boolean; +} + +export class BridgeProtocolError extends Error { + constructor( + message: string, + public readonly status?: number, + ) { + super(message); + this.name = 'BridgeProtocolError'; + } +} + +export function bridgeWorkerPath(workerId: string): string { + return `/bridge/workers/${encodeURIComponent(workerId)}`; +} diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts new file mode 100644 index 00000000..30f07993 --- /dev/null +++ b/packages/code/src/worker.test.ts @@ -0,0 +1,539 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import { BridgeWorker, BridgeWorkspaceQuarantinedError } from './worker.js'; + +import type { BridgeAssignment } from './protocol.js'; + +test('worker forwards a fenced assignment to the sandbox and settles the result', async () => { + const requests: Array<{ url: string; init?: RequestInit }> = []; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + requests.push({ url, init }); + if (url.endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1/', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2/', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + const assignment: BridgeAssignment = { + protocolVersion: 1, + assignmentId: 'assignment-1', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 3, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { + body: { language: 'bash' }, + headers: { 'X-Execution-Manifest': 'signed' }, + }, + }; + + await worker.executeAndSettle(assignment); + + assert.equal(requests.length, 2); + assert.equal( + requests[0].url, + 'http://127.0.0.1:2000/sessions/rt-user-1/api/v2/execute', + ); + assert.equal( + (requests[0].init?.headers as Record)[ + 'X-Runtime-Session-Id' + ], + 'rt-user-1', + ); + assert.match(requests[1].url, /assignments\/assignment-1\/settle$/); + assert.deepEqual(JSON.parse(String(requests[1].init?.body)), { + protocolVersion: 1, + generation: 3, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + incarnationId: 'incarnation-00000001', + status: 'fulfilled', + result: { session_id: 'run-1', files: [] }, + }); +}); + +test('worker aborts sandbox execution at the absolute assignment deadline', async () => { + let settlement: Record | undefined; + const fetchImpl: typeof fetch = async (input, init) => { + if (String(input).endsWith('/execute')) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + settlement = JSON.parse(String(init?.body)) as Record; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-deadline', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 30).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.incarnationId, 'incarnation-00000001'); +}); + +test('worker refreshes its registration during a long assignment', async () => { + let registrations = 0; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + registrations += 1; + return new Response( + JSON.stringify({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + registeredAt: new Date().toISOString(), + leaseTtlMs: 50, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + if (url.endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 90)); + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response( + JSON.stringify({ + protocolVersion: 1, + accepted: true, + body: init?.body, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + await worker.register(); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-heartbeat', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.ok(registrations >= 2); +}); + +test('worker routes a hintless assignment to an ephemeral template session', async () => { + let executeUrl = ''; + let runtimeSessionHeader = ''; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + executeUrl = url; + runtimeSessionHeader = (init?.headers as Record)[ + 'X-Runtime-Session-Id' + ]; + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'hintless-assignment', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal( + executeUrl, + 'http://127.0.0.1:2000/sessions/assignment-hintless-assignment/api/v2/execute', + ); + assert.equal(runtimeSessionHeader, 'assignment-hintless-assignment'); +}); + +test('worker quarantines a fulfilled stateful settlement rejected by Code API', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response(JSON.stringify({ error: 'assignment was fenced' }), { + status: 409, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'fenced-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); +}); + +test('worker surfaces a definite stateless settlement rejection directly', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response(JSON.stringify({ error: 'assignment was fenced' }), { + status: 409, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'fenced-stateless-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }), + (error: unknown) => + error instanceof Error && + error.name === 'BridgeProtocolError' && + error.message === 'assignment was fenced', + ); +}); + +test('worker retries an ambiguous settlement before the deadline', async () => { + let settlementAttempts = 0; + const fetchImpl: typeof fetch = async (input) => { + const url = String(input); + if (url.endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + settlementAttempts += 1; + if (settlementAttempts === 1) throw new TypeError('connection reset'); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'retry-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(settlementAttempts, 2); +}); + +test('worker quarantines stateful reuse after settlement stays ambiguous', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + throw new TypeError('connection reset'); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'ambiguous-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 50).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); +}); + +test('worker quarantines a stateful workspace after the sandbox request aborts', async () => { + let settlementAttempted = false; + const fetchImpl: typeof fetch = async (input, init) => { + if (String(input).endsWith('/execute')) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + settlementAttempted = true; + return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'aborted-execution', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 30).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); + assert.equal(settlementAttempted, false); +}); + +test('worker surfaces quarantine when shutdown aborts stateful execution', async () => { + const controller = new AbortController(); + let executeStarted = false; + const assignment: BridgeAssignment = { + protocolVersion: 1, + assignmentId: 'shutdown-execution', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + return new Response( + JSON.stringify({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + if (url.endsWith('/lease')) { + return new Response( + JSON.stringify({ protocolVersion: 1, assignment }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + if (url.endsWith('/execute')) { + executeStarted = true; + setTimeout(() => controller.abort(), 10); + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.run(controller.signal), + BridgeWorkspaceQuarantinedError, + ); + assert.equal(executeStarted, true); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts new file mode 100644 index 00000000..263c58c5 --- /dev/null +++ b/packages/code/src/worker.ts @@ -0,0 +1,444 @@ +import { randomBytes } from 'node:crypto'; + +import { + BRIDGE_PROTOCOL_VERSION, + BridgeProtocolError, + bridgeWorkerPath, +} from './protocol.js'; + +import type { + BridgeAssignment, + BridgeLeaseResponse, + BridgeSettlement, + BridgeSettlementResponse, + BridgeWorkerCapabilities, + BridgeWorkerRegistrationResponse, +} from './protocol.js'; + +export interface BridgeWorkerOptions { + codeApiUrl: string; + token: string; + workerId: string; + sandboxEndpoint: string; + capabilities: BridgeWorkerCapabilities; + leaseWaitMs?: number; + reconnectDelayMs?: number; + fetchImpl?: typeof fetch; + onError?: (error: unknown) => void; + incarnationId?: string; +} + +const DEFAULT_LEASE_WAIT_MS = 25_000; +const DEFAULT_RECONNECT_DELAY_MS = 1_000; +const DEFAULT_REGISTRATION_TTL_MS = 60_000; +const MIN_REGISTRATION_HEARTBEAT_MS = 25; +const SETTLEMENT_RETRY_DELAY_MS = 100; +const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; + +function normalizedBaseUrl(value: string): string { + return value.replace(/\/+$/, ''); +} + +function errorMessage(value: object): string | undefined { + if ('error' in value && typeof value.error === 'string') return value.error; + return undefined; +} + +export class BridgeWorkspaceQuarantinedError extends Error { + constructor( + message: string, + public readonly cause?: unknown, + ) { + super(message); + this.name = 'BridgeWorkspaceQuarantinedError'; + } +} + +export class BridgeWorker { + private readonly fetchImpl: typeof fetch; + private readonly codeApiUrl: string; + private readonly sandboxEndpoint: string; + private readonly incarnationId: string; + private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; + + constructor(private readonly options: BridgeWorkerOptions) { + this.fetchImpl = options.fetchImpl ?? fetch; + this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); + this.sandboxEndpoint = normalizedBaseUrl(options.sandboxEndpoint); + this.incarnationId = + options.incarnationId ?? randomBytes(18).toString('base64url'); + } + + async register( + signal?: AbortSignal, + ): Promise { + const registration = await this.request( + `${this.codeApiUrl}/bridge/workers/register`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: this.options.workerId, + incarnationId: this.incarnationId, + capabilities: this.options.capabilities, + }, + signal, + ); + if (registration.incarnationId !== this.incarnationId) { + throw new BridgeProtocolError( + 'Code API registered a different worker incarnation', + ); + } + this.registrationTtlMs = registration.leaseTtlMs; + return registration; + } + + async lease(signal?: AbortSignal): Promise { + const response = await this.request( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/lease`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + waitMs: this.options.leaseWaitMs ?? DEFAULT_LEASE_WAIT_MS, + incarnationId: this.incarnationId, + }, + signal, + ); + if ( + response.assignment != null && + response.assignment.incarnationId !== this.incarnationId + ) { + throw new BridgeProtocolError( + 'Code API leased an assignment for a different worker incarnation', + ); + } + return response.assignment; + } + + async run(signal?: AbortSignal): Promise { + while (!signal?.aborted) { + try { + await this.register(signal); + const assignment = await this.lease(signal); + if (!assignment) continue; + await this.executeAndSettle(assignment, signal); + } catch (error) { + if (error instanceof BridgeWorkspaceQuarantinedError) { + throw error; + } + if (signal?.aborted) return; + if ( + error instanceof BridgeProtocolError && + (error.status === 401 || error.status === 403 || error.status === 409) + ) { + throw error; + } + this.options.onError?.(error); + const delay = + this.options.reconnectDelayMs ?? DEFAULT_RECONNECT_DELAY_MS; + await new Promise((resolve) => setTimeout(resolve, delay)); + } + } + } + + async executeAndSettle( + assignment: BridgeAssignment, + signal?: AbortSignal, + ): Promise { + const executionController = new AbortController(); + const abortExecution = (): void => executionController.abort(); + signal?.addEventListener('abort', abortExecution, { once: true }); + const deadlineDelay = Math.max( + 0, + Date.parse(assignment.expiresAt) - Date.now(), + ); + const deadlineTimer = setTimeout( + () => executionController.abort(), + deadlineDelay, + ); + const heartbeatController = new AbortController(); + let heartbeatError: unknown; + const heartbeat = this.maintainRegistration( + heartbeatController.signal, + executionController, + ).catch((error) => { + heartbeatError = error; + executionController.abort(); + }); + const cancellationController = new AbortController(); + const cancellationWatcher = this.watchCancellation( + assignment, + executionController, + cancellationController.signal, + ); + let settlement: BridgeSettlement; + let ambiguousSandboxError: unknown; + let sandboxRejectedExecution = false; + try { + const sandboxSessionId = this.sandboxSessionIdFor(assignment); + const headers = { + ...assignment.request.headers, + ...(sandboxSessionId + ? { 'X-Runtime-Session-Id': sandboxSessionId } + : {}), + }; + const response = await this.fetchImpl( + `${this.sandboxEndpointFor(assignment)}/execute`, + { + method: 'POST', + headers: { + ...headers, + 'Content-Type': 'application/json', + }, + body: JSON.stringify(assignment.request.body), + signal: executionController.signal, + }, + ); + const payload = (await response.json()) as object; + if (heartbeatError != null) throw heartbeatError; + if (!response.ok) { + sandboxRejectedExecution = true; + throw new BridgeProtocolError( + errorMessage(payload) ?? + `Sandbox rejected execution with HTTP ${response.status}`, + response.status, + ); + } + settlement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, + status: 'fulfilled', + result: payload, + }; + } catch (error) { + if ( + assignment.runtimeSessionId != null && + !sandboxRejectedExecution + ) { + ambiguousSandboxError = error; + } + settlement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, + status: 'rejected', + error: + error instanceof Error ? error.message : 'Sandbox execution failed', + }; + } + + clearTimeout(deadlineTimer); + cancellationController.abort(); + await cancellationWatcher; + try { + if (ambiguousSandboxError != null) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined after an ambiguous sandbox execution`, + ambiguousSandboxError, + ); + } + await this.settleWithRetry(assignment, settlement, signal); + } finally { + heartbeatController.abort(); + await heartbeat; + signal?.removeEventListener('abort', abortExecution); + } + } + + private sandboxSessionIdFor( + assignment: BridgeAssignment, + ): string | undefined { + if (assignment.runtimeSessionId != null) { + return assignment.runtimeSessionId; + } + if (this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER)) { + return `assignment-${assignment.assignmentId}`; + } + return undefined; + } + + private sandboxEndpointFor(assignment: BridgeAssignment): string { + if (assignment.runtimeSessionId == null) { + if (!this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER)) { + return this.sandboxEndpoint; + } + return this.sandboxEndpoint.replace( + RUNTIME_SESSION_PLACEHOLDER, + encodeURIComponent(`assignment-${assignment.assignmentId}`), + ); + } + if ( + this.options.capabilities.statefulWorkspace !== true || + !this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER) + ) { + throw new BridgeProtocolError( + 'Stateful assignments require a sandbox endpoint template containing {runtimeSessionId}', + ); + } + return this.sandboxEndpoint.replace( + RUNTIME_SESSION_PLACEHOLDER, + encodeURIComponent(assignment.runtimeSessionId), + ); + } + + private async maintainRegistration( + signal: AbortSignal, + executionController: AbortController, + ): Promise { + while (!signal.aborted && !executionController.signal.aborted) { + await this.delay( + Math.max( + MIN_REGISTRATION_HEARTBEAT_MS, + Math.floor(this.registrationTtlMs / 2), + ), + signal, + ); + if (signal.aborted || executionController.signal.aborted) return; + await this.register(signal); + } + } + + private async delay(ms: number, signal: AbortSignal): Promise { + if (signal.aborted) return; + await new Promise((resolve) => { + const onAbort = (): void => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + signal.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal.addEventListener('abort', onAbort, { once: true }); + }); + } + + private assignmentUrl(assignment: BridgeAssignment, action: string): string { + return ( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + + `/assignments/${encodeURIComponent(assignment.assignmentId)}/${action}` + ); + } + + private async settleWithRetry( + assignment: BridgeAssignment, + settlement: BridgeSettlement, + signal?: AbortSignal, + ): Promise { + const deadlineAtMs = Date.parse(assignment.expiresAt); + const settlementController = new AbortController(); + const abortSettlement = (): void => settlementController.abort(); + signal?.addEventListener('abort', abortSettlement, { once: true }); + const deadlineTimer = setTimeout( + () => settlementController.abort(), + Math.max(0, deadlineAtMs - Date.now()), + ); + let lastError: unknown; + try { + while (!settlementController.signal.aborted) { + try { + await this.request( + this.assignmentUrl(assignment, 'settle'), + settlement, + settlementController.signal, + ); + return; + } catch (error) { + lastError = error; + if (signal?.aborted) break; + if ( + error instanceof BridgeProtocolError && + error.status != null && + error.status < 500 && + error.status !== 408 && + error.status !== 429 + ) { + if ( + assignment.runtimeSessionId != null && + settlement.status === 'fulfilled' + ) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined after Code API rejected its fulfilled settlement`, + error, + ); + } + throw error; + } + const remainingMs = deadlineAtMs - Date.now(); + if (remainingMs <= 0) break; + await this.delay( + Math.min(SETTLEMENT_RETRY_DELAY_MS, remainingMs), + settlementController.signal, + ); + } + } + } finally { + clearTimeout(deadlineTimer); + signal?.removeEventListener('abort', abortSettlement); + } + if (assignment.runtimeSessionId != null) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined after ambiguous settlement delivery`, + lastError, + ); + } + if (lastError instanceof Error) throw lastError; + throw new BridgeProtocolError('Bridge settlement deadline expired'); + } + + private async watchCancellation( + assignment: BridgeAssignment, + executionController: AbortController, + signal: AbortSignal, + ): Promise { + while (!signal.aborted && !executionController.signal.aborted) { + await new Promise((resolve) => setTimeout(resolve, 500)); + if (signal.aborted || executionController.signal.aborted) return; + try { + const response = await this.request<{ cancelled: boolean }>( + this.assignmentUrl(assignment, 'cancellation'), + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + }, + signal, + ); + if (response.cancelled) { + executionController.abort(); + return; + } + } catch (error) { + if (signal.aborted) return; + if (error instanceof BridgeProtocolError && error.status === 404) { + executionController.abort(); + return; + } + } + } + } + + private async request( + url: string, + body: object, + signal?: AbortSignal, + ): Promise { + const response = await this.fetchImpl(url, { + method: 'POST', + headers: { + Authorization: `Bearer ${this.options.token}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify(body), + signal, + }); + const payload = (await response.json()) as object; + if (!response.ok) { + throw new BridgeProtocolError( + errorMessage(payload) ?? + `Bridge request failed with HTTP ${response.status}`, + response.status, + ); + } + return payload as T; + } +} diff --git a/packages/code/tsconfig.json b/packages/code/tsconfig.json new file mode 100644 index 00000000..6c7ea362 --- /dev/null +++ b/packages/code/tsconfig.json @@ -0,0 +1,15 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "NodeNext", + "moduleResolution": "NodeNext", + "strict": true, + "declaration": true, + "sourceMap": true, + "outDir": "dist", + "rootDir": "src", + "skipLibCheck": true, + "types": ["node"] + }, + "include": ["src/**/*.ts"] +} diff --git a/service/Dockerfile b/service/Dockerfile index 762790dc..00680111 100644 --- a/service/Dockerfile +++ b/service/Dockerfile @@ -18,6 +18,7 @@ COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY service/scripts ./scripts COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ RUN bun build ./src/file-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' RUN bun build ./src/api-server.ts --minify --outdir .build-api --target bun --external '@opentelemetry/*' @@ -66,6 +67,7 @@ ENV NODE_ENV=development COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ EXPOSE 3000 9230 CMD ["bun", "run", "--watch", "src/file-server.ts"] diff --git a/service/Dockerfile.api b/service/Dockerfile.api index 419bdc95..f1fdf9c8 100644 --- a/service/Dockerfile.api +++ b/service/Dockerfile.api @@ -18,6 +18,7 @@ RUN cd /temp/prod && bun install --frozen-lockfile --production FROM base AS builder COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src +COPY packages/code/src /packages/code/src COPY service/scripts ./scripts COPY shared /shared COPY service/tsconfig.json ./ @@ -46,6 +47,7 @@ FROM base AS development ENV NODE_ENV=development COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ EXPOSE 3112 9230 diff --git a/service/Dockerfile.local b/service/Dockerfile.local index f932deee..cbb7af13 100644 --- a/service/Dockerfile.local +++ b/service/Dockerfile.local @@ -17,6 +17,7 @@ FROM base AS builder COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ RUN bun build ./src/local-api.ts --minify --outdir .build --target bun --external '@opentelemetry/*' @@ -35,5 +36,6 @@ ENV NODE_ENV=development COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ CMD ["bun", "run", "--watch", "src/local-api.ts"] diff --git a/service/Dockerfile.worker b/service/Dockerfile.worker index 9d1a4322..e99c16c4 100644 --- a/service/Dockerfile.worker +++ b/service/Dockerfile.worker @@ -19,6 +19,7 @@ RUN cd /temp/prod && bun install --frozen-lockfile --production FROM base AS builder COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ RUN bun build ./src/worker-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' @@ -43,6 +44,7 @@ FROM base AS development ENV NODE_ENV=development COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ EXPOSE 3113 9230 diff --git a/service/rollup.config.js b/service/rollup.config.js index 2400f72d..0e7a8d8b 100644 --- a/service/rollup.config.js +++ b/service/rollup.config.js @@ -38,7 +38,11 @@ export default { commonjs(), typescript({ tsconfig: './tsconfig.esm.json', - include: ['src/**/*.ts', '../shared/telemetry-core.ts'], + include: [ + 'src/**/*.ts', + '../shared/telemetry-core.ts', + '../packages/code/src/protocol.ts', + ], sourceMap: true, declaration: false, declarationMap: false, diff --git a/service/src/api-server.ts b/service/src/api-server.ts index 78689826..8578460b 100644 --- a/service/src/api-server.ts +++ b/service/src/api-server.ts @@ -18,6 +18,7 @@ import { requestErrorLogger, requestNotFoundLogger } from './middleware/request- import { localAuth } from './auth/local'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; +import bridgeRouter from './bridge/router'; import { connection } from './queue'; import { metricsHandler } from './metrics'; import { httpMetricsMiddleware } from './middleware/httpMetrics'; @@ -51,6 +52,7 @@ app.get('/v1/health', async (_, res) => { } }); +v1.use('/bridge', bridgeRouter); v1.use(isLocalMode ? localAuth : apiKeyAuth); v1.use(serviceRouter); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts new file mode 100644 index 00000000..8653bade --- /dev/null +++ b/service/src/bridge/router.ts @@ -0,0 +1,248 @@ +import { timingSafeEqual } from 'crypto'; + +import { Router } from 'express'; +import type { NextFunction, Request, RequestHandler, Response } from 'express'; +import type { BridgeWorkerRegistration } from '../../../packages/code/src/protocol'; +import type { CodeBridgeSettlement } from './store'; + +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { connection } from '../queue'; +import { env } from '../config'; +import { BridgeStoreError, RedisBridgeStore } from './store'; + +const WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; +const INCARNATION_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; +const MAX_LEASE_WAIT_MS = 30_000; + +export const bridgeStore = new RedisBridgeStore(connection); + +function sameToken(left: string, right: string): boolean { + const leftBuffer = Buffer.from(left); + const rightBuffer = Buffer.from(right); + return ( + leftBuffer.length === rightBuffer.length && + timingSafeEqual(leftBuffer, rightBuffer) + ); +} + +function bridgeAuth(req: Request, res: Response, next: NextFunction): void { + if (!env.BRIDGE_TOKEN) { + res.status(503).json({ error: 'Code bridge is not configured' }); + return; + } + const token = + req + .header('Authorization') + ?.match(/^Bearer\s+(.+)$/i)?.[1] + ?.trim() ?? ''; + if (!token || !sameToken(token, env.BRIDGE_TOKEN)) { + res.status(401).json({ error: 'Invalid code bridge worker token' }); + return; + } + next(); +} + +function validWorkerId(value: string): boolean { + return WORKER_ID_PATTERN.test(value); +} + +function validIncarnationId(value: unknown): value is string { + return typeof value === 'string' && INCARNATION_ID_PATTERN.test(value); +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function asyncRoute( + handler: (req: Request, res: Response) => Promise, +): RequestHandler { + return (req, res, next) => { + void handler(req, res).catch(next); + }; +} + +function sendStoreError(error: BridgeStoreError, res: Response): void { + const status = + error.code === 'ASSIGNMENT_NOT_FOUND' + ? 404 + : error.code === 'WORKER_BUSY' + ? 503 + : 409; + res.status(status).json({ error: error.message, code: error.code }); +} + +function isSettlement(value: unknown): value is CodeBridgeSettlement { + if (!isRecord(value)) return false; + if ( + value.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof value.generation !== 'number' || + !Number.isSafeInteger(value.generation) || + value.generation < 1 || + typeof value.leaseToken !== 'string' || + value.leaseToken.length < 32 || + !validIncarnationId(value.incarnationId) + ) { + return false; + } + if (value.status === 'rejected') { + return typeof value.error === 'string' && value.error.length <= 4096; + } + return ( + value.status === 'fulfilled' && + typeof value.result === 'object' && + value.result !== null + ); +} + +const router = Router(); +router.use(bridgeAuth); + +router.post( + '/workers/register', + asyncRoute(async (req, res) => { + const registration = req.body as unknown; + if ( + !isRecord(registration) || + registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof registration.workerId !== 'string' || + !validWorkerId(registration.workerId) || + !validIncarnationId(registration.incarnationId) || + !isRecord(registration.capabilities) || + typeof registration.capabilities.statefulWorkspace !== 'boolean' || + typeof registration.capabilities.sandboxProfile !== 'string' || + registration.capabilities.sandboxProfile.trim().length === 0 || + registration.capabilities.sandboxProfile.length > 128 || + !Array.isArray(registration.capabilities.runtimes) || + registration.capabilities.runtimes.length > 32 || + !registration.capabilities.runtimes.every( + (runtime) => + typeof runtime === 'string' && + runtime.length > 0 && + runtime.length <= 64, + ) || + (registration.capabilities.policyDigest !== undefined && + (typeof registration.capabilities.policyDigest !== 'string' || + !/^[a-f0-9]{64}$/.test(registration.capabilities.policyDigest))) + ) { + res.status(400).json({ error: 'Invalid bridge worker registration' }); + return; + } + if ( + env.BRIDGE_WORKER_ID && + registration.workerId !== env.BRIDGE_WORKER_ID + ) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { + await bridgeStore.register( + registration as unknown as BridgeWorkerRegistration, + ); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }), +); + +router.post( + '/workers/:workerId/lease', + asyncRoute(async (req, res) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + const requestedWait = Number(body.waitMs ?? 25_000); + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isFinite(requestedWait) || + requestedWait < 0 + ) { + res.status(400).json({ error: 'Invalid bridge lease request' }); + return; + } + if (env.BRIDGE_WORKER_ID && workerId !== env.BRIDGE_WORKER_ID) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { + const assignment = await bridgeStore.lease( + workerId, + body.incarnationId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/assignments/:assignmentId/settle', + asyncRoute(async (req, res) => { + const settlement = req.body as unknown; + if (!isSettlement(settlement)) { + res.status(400).json({ error: 'Invalid bridge settlement' }); + return; + } + try { + await bridgeStore.settle( + req.params.workerId, + req.params.assignmentId, + settlement, + ); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + accepted: true, + }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/assignments/:assignmentId/cancellation', + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) + ) { + res.status(400).json({ error: 'Invalid bridge cancellation request' }); + return; + } + const cancelled = await bridgeStore.cancelled( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); + }), +); + +export default router; diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts new file mode 100644 index 00000000..6a4cf77b --- /dev/null +++ b/service/src/bridge/store.test.ts @@ -0,0 +1,659 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { getEventListeners } from 'node:events'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import type * as t from '../types'; +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { RedisBridgeStore } from './store'; + +const redis = new RedisMock() as unknown as Redis; +const store = new RedisBridgeStore(redis); +const incarnationId = 'incarnation-00000001'; +const redisEval = redis.eval.bind(redis); +const redisDel = redis.del.bind(redis); + +afterEach(async () => { + redis.eval = redisEval as Redis['eval']; + redis.del = redisDel as Redis['del']; + await redis.flushall(); +}); + +describe('RedisBridgeStore', () => { + test('delivers and settles one fenced stateful assignment', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: { 'X-Execution-Manifest': 'signed' }, + runtimeSessionId: 'rt-user-1', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('vm-1', incarnationId, 1_000); + expect(assignment).toBeDefined(); + expect(assignment?.runtimeSessionId).toBe('rt-user-1'); + + await store.settle('vm-1', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + language: 'bash', + version: '5.2.0', + session_id: 'run-1', + files: [], + }, + }); + + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'run-1' }, + }); + }); + + test('rejects dispatch to an offline worker', async () => { + const controller = new AbortController(); + await expect( + store.dispatch({ + workerId: 'offline', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: controller.signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_OFFLINE' }); + }); + + test('rejects a stale lease token', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('vm-1', incarnationId, 1_000); + + await expect( + store.settle('vm-1', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: 'stale-token-that-is-long-enough-to-pass-validation', + incarnationId, + status: 'rejected', + error: 'unused', + }), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('fences a replaced worker incarnation', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_FENCED' }); + }); + + test('a stale incarnation poll cannot consume replacement work', async () => { + const replacementIncarnationId = 'incarnation-00000002'; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'restarted-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const stalePoll = store.lease('restarted-worker', incarnationId, 100); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'restarted-worker', + incarnationId: replacementIncarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'restarted-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect(stalePoll).resolves.toBeUndefined(); + await expect( + store.lease('restarted-worker', replacementIncarnationId, 1_000), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('dispatch retries atomically against a replacement incarnation', async () => { + const workerId = 'racing-worker'; + const replacementIncarnationId = 'incarnation-00000002'; + const capabilities = { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [] as string[], + }; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities, + }); + const originalEval = redis.eval.bind(redis); + let replaced = false; + redis.eval = (async (...args: Parameters) => { + if (!replaced && String(args[0]).includes("redis.call('RPUSH'")) { + replaced = true; + const replacement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId: replacementIncarnationId, + capabilities, + }; + await redis.set( + `codeapi:bridge:v1:worker:${workerId}`, + JSON.stringify(replacement), + 'EX', + 60, + ); + await redis.set( + `codeapi:bridge:v1:worker:${workerId}:incarnation`, + replacementIncarnationId, + 'EX', + 60, + ); + } + return originalEval(...args); + }) as Redis['eval']; + const controller = new AbortController(); + const completion = store.dispatch({ + workerId, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + const assignment = await store.lease( + workerId, + replacementIncarnationId, + 1_000, + ); + expect(assignment?.incarnationId).toBe(replacementIncarnationId); + redis.eval = originalEval as Redis['eval']; + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('defers worker replacement while an assignment is active', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'busy-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'busy-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('busy-worker', incarnationId, 1_000); + expect(assignment).toBeDefined(); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'busy-worker', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'busy-worker', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).resolves.toBeUndefined(); + }); + + test('recovers only the assignment owner after registration expiry', async () => { + const workerId = 'expired-registration-worker'; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease(workerId, incarnationId, 1_000); + expect(assignment).toBeDefined(); + await redis.del( + `codeapi:bridge:v1:worker:${workerId}`, + `codeapi:bridge:v1:worker:${workerId}:incarnation`, + ); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).resolves.toBeUndefined(); + + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('removes abort listeners after each settlement poll delay', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'listener-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'listener-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 350, + signal: controller.signal, + }); + + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + expect(getEventListeners(controller.signal, 'abort')).toHaveLength(0); + }); + + test('keeps assignment state through deadlines longer than ten minutes', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'long-running-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'long-running-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 15 * 60_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + const [assignmentKey] = await redis.keys('codeapi:bridge:v1:assignment:*'); + + expect(await redis.ttl(assignmentKey)).toBeGreaterThan(10 * 60); + expect( + await redis.pttl('codeapi:bridge:v1:worker:long-running-worker:lock'), + ).toBeGreaterThan(10 * 60_000); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('observes a settlement accepted during the final poll delay', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'deadline-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const deadlineAtMs = Date.now() + 500; + const completion = store.dispatch({ + workerId: 'deadline-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-deadline', + deadlineAtMs, + signal: controller.signal, + }); + const assignment = await store.lease( + 'deadline-worker', + incarnationId, + 1_000, + ); + expect(assignment).toBeDefined(); + await new Promise((resolve) => + setTimeout(resolve, Math.max(0, deadlineAtMs - Date.now() - 30)), + ); + await store.settle('deadline-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + language: 'bash', + version: '5.2.0', + session_id: 'run-deadline', + files: [], + }, + }); + + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'run-deadline' }, + }); + }); + + test('preserves a committed result across transient cleanup failures', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'cleanup-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'cleanup-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-cleanup', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'cleanup-worker', + incarnationId, + 1_000, + ); + const originalDel = redis.del.bind(redis); + let cleanupAttempts = 0; + redis.del = (async (...args: Parameters) => { + cleanupAttempts += 1; + if (cleanupAttempts === 1) throw new Error('transient cleanup failure'); + return originalDel(...args); + }) as Redis['del']; + await store.settle('cleanup-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + language: 'bash', + version: '5.2.0', + session_id: 'run-cleanup', + files: [], + }, + }); + + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'run-cleanup' }, + }); + expect(cleanupAttempts).toBeGreaterThanOrEqual(2); + redis.del = originalDel as Redis['del']; + }); + + test('releases the worker lock when generation allocation fails', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const originalIncr = redis.incr.bind(redis); + let failOnce = true; + redis.incr = (async (...args: Parameters) => { + if (failOnce) { + failOnce = false; + throw new Error('incr failed'); + } + return originalIncr(...args); + }) as Redis['incr']; + const controller = new AbortController(); + await expect( + store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: controller.signal, + }), + ).rejects.toThrow('incr failed'); + redis.incr = originalIncr as Redis['incr']; + + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: controller.signal, + }); + const assignment = await store.lease('vm-1', incarnationId, 500); + expect(assignment).toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('quarantines a workspace when result finalization fails', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-user-1', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + finalize: async () => { + await expect( + store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-user-1', + deadlineAtMs: Date.now() + 1_000, + signal: controller.signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + throw new Error('restore failed'); + }, + }); + const assignment = await store.lease('vm-1', incarnationId, 1_000); + await store.settle('vm-1', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + language: 'bash', + version: '5.2.0', + session_id: 'run-1', + files: [], + }, + }); + + await expect(completion).rejects.toThrow('restore failed'); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_QUARANTINED' }); + + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + await expect( + store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-user-1', + deadlineAtMs: Date.now() + 1_000, + signal: controller.signal, + }), + ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + }); +}); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts new file mode 100644 index 00000000..e0b09b8a --- /dev/null +++ b/service/src/bridge/store.ts @@ -0,0 +1,629 @@ +import { createHash, randomBytes } from 'crypto'; + +import type Redis from 'ioredis'; +import type * as t from '../types'; +import type { + BridgeAssignment, + BridgeSettlement, + BridgeWorkerRegistration, +} from '../../../packages/code/src/protocol'; + +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; + +const PREFIX = 'codeapi:bridge:v1'; +const POLL_INTERVAL_MS = 100; +const DEFAULT_WORKER_TTL_SECONDS = 60; + +export type CodeBridgeAssignment = BridgeAssignment; +export type CodeBridgeSettlement = BridgeSettlement< + t.ExecuteResponse & { + session_id: string; + files?: t.FileRefs; + run?: t.ExecuteResponse['run']; + } +>; + +export class BridgeStoreError extends Error { + constructor( + public readonly code: + | 'WORKER_OFFLINE' + | 'WORKER_BUSY' + | 'ASSIGNMENT_EXPIRED' + | 'ASSIGNMENT_FENCED' + | 'ASSIGNMENT_NOT_FOUND' + | 'WORKER_FENCED' + | 'WORKER_QUARANTINED' + | 'WORKSPACE_QUARANTINED' + | 'WORKER_MISMATCH', + message: string, + ) { + super(message); + this.name = 'BridgeStoreError'; + } +} + +interface StoredAssignment extends CodeBridgeAssignment { + leaseTokenHash: string; +} + +function workerKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}`; +} + +function workerIncarnationKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation`; +} + +function incarnationFenceKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:fenced`; +} + +function quarantineKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:quarantined`; +} + +function workspaceQuarantineKey( + workerId: string, + runtimeSessionId: string, +): string { + const sessionHash = createHash('sha256') + .update(runtimeSessionId) + .digest('hex'); + return `${PREFIX}:worker:${workerId}:workspace:${sessionHash}:quarantined`; +} + +function queueKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:assignments`; +} + +function generationKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:generation`; +} + +function lockKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:lock`; +} + +function lockIncarnationKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:lock:incarnation`; +} + +function assignmentKey(assignmentId: string): string { + return `${PREFIX}:assignment:${assignmentId}`; +} + +function settlementKey(assignmentId: string): string { + return `${PREFIX}:assignment:${assignmentId}:settlement`; +} + +function cancellationKey(assignmentId: string): string { + return `${PREFIX}:assignment:${assignmentId}:cancelled`; +} + +function tokenHash(token: string): string { + return createHash('sha256').update(token).digest('hex'); +} + +function assignmentTtlSeconds(deadlineAtMs: number): number { + return Math.max(1, Math.ceil((deadlineAtMs - Date.now()) / 1000) + 30); +} + +async function delay(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted === true) return; + await new Promise((resolve) => { + const onAbort = (): void => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} + +export class RedisBridgeStore { + constructor( + private readonly redis: Redis, + private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, + ) {} + + async register(registration: BridgeWorkerRegistration): Promise { + const script = [ + 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', + 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', + 'local current = redis.call(\'GET\', KEYS[4])', + 'if not current and redis.call(\'EXISTS\', KEYS[5]) == 1 then', + ' local owner = redis.call(\'GET\', KEYS[6])', + ' if owner ~= ARGV[1] then return -3 end', + 'end', + 'if current then', + ' if current ~= ARGV[1] then', + ' if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -3 end', + ' redis.call(\'SET\', ARGV[4] .. current .. \':fenced\', \"1\")', + ' end', + 'end', + 'redis.call(\'SET\', KEYS[1], ARGV[2], \"EX\", ARGV[3])', + 'redis.call(\'SET\', KEYS[4], ARGV[1], \"EX\", ARGV[3])', + 'return 1', + ].join('\n'); + const result = Number( + await this.redis.eval( + script, + 6, + workerKey(registration.workerId), + incarnationFenceKey(registration.workerId, registration.incarnationId), + quarantineKey(registration.workerId, registration.incarnationId), + workerIncarnationKey(registration.workerId), + lockKey(registration.workerId), + lockIncarnationKey(registration.workerId), + registration.incarnationId, + JSON.stringify(registration), + String(this.workerTtlSeconds), + `${PREFIX}:worker:${registration.workerId}:incarnation:`, + ), + ); + if (result === -2) { + throw new BridgeStoreError( + 'WORKER_QUARANTINED', + 'Bridge worker incarnation is quarantined', + ); + } + if (result === -1) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Bridge worker incarnation was replaced', + ); + } + if (result === -3) { + throw new BridgeStoreError( + 'WORKER_BUSY', + 'Bridge worker cannot be replaced during an active assignment', + ); + } + } + + async dispatch(args: { + workerId: string; + body: t.PayloadBody; + headers: Record; + runtimeSessionId?: string; + deadlineAtMs: number; + signal: AbortSignal; + finalize?: ( + settlement: CodeBridgeSettlement, + ) => Promise; + }): Promise { + let registration = await this.registration(args.workerId); + if (registration == null) { + throw new BridgeStoreError( + 'WORKER_OFFLINE', + `Bridge worker ${args.workerId} is offline`, + ); + } + if ( + args.runtimeSessionId !== undefined && + registration.capabilities.statefulWorkspace !== true + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} does not provide a stateful workspace`, + ); + } + if ( + args.runtimeSessionId !== undefined && + (await this.redis.exists( + workspaceQuarantineKey(args.workerId, args.runtimeSessionId), + )) === 1 + ) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace is quarantined after an incomplete result commit', + ); + } + + const assignmentId = randomBytes(18).toString('base64url'); + const leaseToken = randomBytes(32).toString('base64url'); + const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs); + const locked = await this.acquireLock( + args.workerId, + assignmentId, + registration.incarnationId, + ttlSeconds, + ); + if (!locked) { + throw new BridgeStoreError( + 'WORKER_BUSY', + `Bridge worker ${args.workerId} is busy`, + ); + } + + let assignment: StoredAssignment | undefined; + let resultCommitted = false; + try { + const generation = await this.redis.incr(generationKey(args.workerId)); + assignment = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + assignmentId, + workerId: args.workerId, + incarnationId: registration.incarnationId, + generation, + leaseToken, + leaseTokenHash: tokenHash(leaseToken), + expiresAt: new Date(args.deadlineAtMs).toISOString(), + runtimeSessionId: args.runtimeSessionId, + request: { + body: args.body, + headers: args.headers, + }, + }; + let queued = false; + for (let attempt = 0; attempt < 8 && !queued; attempt += 1) { + assignment.incarnationId = registration.incarnationId; + queued = await this.enqueueForActiveIncarnation(assignment, ttlSeconds); + if (queued) break; + const replacement = await this.registration(args.workerId); + if (replacement == null) { + throw new BridgeStoreError( + 'WORKER_OFFLINE', + `Bridge worker ${args.workerId} went offline during dispatch`, + ); + } + if ( + args.runtimeSessionId !== undefined && + replacement.capabilities.statefulWorkspace !== true + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} does not provide a stateful workspace`, + ); + } + registration = replacement; + } + if (!queued) { + throw new BridgeStoreError( + 'WORKER_OFFLINE', + `Bridge worker ${args.workerId} changed incarnation repeatedly during dispatch`, + ); + } + const settlement = await this.waitForSettlement( + assignment, + args.deadlineAtMs, + args.signal, + ); + if (args.finalize == null) { + resultCommitted = true; + return settlement; + } + try { + const result = await args.finalize(settlement); + resultCommitted = true; + return result; + } catch (error) { + await this.quarantine( + args.workerId, + assignment.incarnationId, + args.runtimeSessionId, + ); + throw error; + } + } finally { + if (resultCommitted) { + try { + await this.cleanupWithRetry(args.workerId, assignmentId, assignment); + } catch { + // The lock and assignment have deadline-derived TTLs. Preserve the + // already committed result rather than turning cleanup availability + // into a client-visible failure that could prompt duplicate work. + } + } else { + await this.cleanupDispatch(args.workerId, assignmentId, assignment); + } + } + } + + async lease( + workerId: string, + incarnationId: string, + waitMs: number, + signal?: AbortSignal, + ): Promise { + const deadline = Date.now() + waitMs; + while (signal?.aborted !== true && Date.now() < deadline) { + const assignmentId = await this.redis.lpop( + queueKey(workerId, incarnationId), + ); + if (assignmentId == null) { + await delay( + Math.min(POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), + signal, + ); + continue; + } + const assignment = await this.readAssignment(assignmentId); + if (assignment == null || assignment.workerId !== workerId) continue; + if (assignment.incarnationId !== incarnationId) continue; + const registration = await this.registration(workerId); + if (registration?.incarnationId !== incarnationId) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Bridge worker incarnation was replaced', + ); + } + if (Date.parse(assignment.expiresAt) <= Date.now()) continue; + const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; + return wireAssignment; + } + return undefined; + } + + async settle( + workerId: string, + assignmentId: string, + settlement: CodeBridgeSettlement, + ): Promise { + const assignment = await this.readAssignment(assignmentId); + if (assignment == null) { + throw new BridgeStoreError( + 'ASSIGNMENT_NOT_FOUND', + 'Bridge assignment was not found', + ); + } + if (assignment.workerId !== workerId) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Bridge assignment belongs to another worker', + ); + } + const registration = await this.registration(workerId); + if ( + settlement.incarnationId !== assignment.incarnationId || + registration?.incarnationId !== settlement.incarnationId || + settlement.generation !== assignment.generation || + tokenHash(settlement.leaseToken) !== assignment.leaseTokenHash + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment lease is stale', + ); + } + if (Date.parse(assignment.expiresAt) <= Date.now()) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment has expired', + ); + } + const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); + const script = [ + 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return 0 end', + 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', + 'return 1', + ].join('\n'); + const accepted = Number( + await this.redis.eval( + script, + 2, + assignmentKey(assignmentId), + settlementKey(assignmentId), + JSON.stringify(settlement), + String(ttlSeconds), + ), + ); + if (accepted !== 1) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment closed before settlement was committed', + ); + } + } + + async cancelled( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { + const assignment = await this.readAssignment(assignmentId); + const registration = await this.registration(workerId); + if ( + assignment == null || + assignment.workerId !== workerId || + assignment.incarnationId !== incarnationId || + registration?.incarnationId !== incarnationId + ) { + return true; + } + return (await this.redis.exists(cancellationKey(assignmentId))) === 1; + } + + async quarantine( + workerId: string, + incarnationId: string, + runtimeSessionId?: string, + ): Promise { + const script = [ + 'redis.call(\'SET\', KEYS[2], \"1\")', + 'if #KEYS == 4 then redis.call(\'SET\', KEYS[4], \"1\") end', + 'local current = redis.call(\'GET\', KEYS[3])', + 'if current == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1], KEYS[3])', + 'end', + 'return 0', + ].join('\n'); + const keys = [ + workerKey(workerId), + quarantineKey(workerId, incarnationId), + workerIncarnationKey(workerId), + ]; + if (runtimeSessionId !== undefined) { + keys.push(workspaceQuarantineKey(workerId, runtimeSessionId)); + } + await this.redis.eval( + script, + keys.length, + ...keys, + incarnationId, + ); + } + + private async registration( + workerId: string, + ): Promise { + const raw = await this.redis.get(workerKey(workerId)); + return raw == null ? undefined : (JSON.parse(raw) as BridgeWorkerRegistration); + } + + private async readAssignment( + assignmentId: string, + ): Promise { + const raw = await this.redis.get(assignmentKey(assignmentId)); + return raw == null ? undefined : (JSON.parse(raw) as StoredAssignment); + } + + private async waitForSettlement( + assignment: StoredAssignment, + deadlineAtMs: number, + signal: AbortSignal, + ): Promise { + while (!signal.aborted && Date.now() < deadlineAtMs) { + const raw = await this.redis.get(settlementKey(assignment.assignmentId)); + if (raw != null) return JSON.parse(raw) as CodeBridgeSettlement; + await delay(POLL_INTERVAL_MS, signal); + } + const closeScript = [ + 'local settlement = redis.call(\'GET\', KEYS[2])', + 'if settlement then return settlement end', + 'redis.call(\'DEL\', KEYS[1])', + 'return nil', + ].join('\n'); + const finalSettlement = await this.redis.eval( + closeScript, + 2, + assignmentKey(assignment.assignmentId), + settlementKey(assignment.assignmentId), + ); + if (finalSettlement != null) { + return JSON.parse(String(finalSettlement)) as CodeBridgeSettlement; + } + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment exceeded its deadline', + ); + } + + private async cancel(assignmentId: string): Promise { + await this.redis.set(cancellationKey(assignmentId), '1', 'EX', 30); + } + + private async enqueueForActiveIncarnation( + assignment: StoredAssignment, + ttlSeconds: number, + ): Promise { + const script = [ + 'if redis.call(\'GET\', KEYS[1]) ~= ARGV[1] then return 0 end', + 'redis.call(\'SET\', KEYS[2], ARGV[2], \"EX\", ARGV[3])', + 'redis.call(\'RPUSH\', KEYS[3], ARGV[4])', + 'redis.call(\'EXPIRE\', KEYS[3], ARGV[3])', + 'redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])', + 'return 1', + ].join('\n'); + const result = await this.redis.eval( + script, + 4, + workerIncarnationKey(assignment.workerId), + assignmentKey(assignment.assignmentId), + queueKey(assignment.workerId, assignment.incarnationId), + lockIncarnationKey(assignment.workerId), + assignment.incarnationId, + JSON.stringify(assignment), + String(ttlSeconds), + assignment.assignmentId, + String(ttlSeconds * 1000), + ); + return Number(result) === 1; + } + + private async acquireLock( + workerId: string, + assignmentId: string, + incarnationId: string, + ttlSeconds: number, + ): Promise { + const script = [ + 'if redis.call(\'EXISTS\', KEYS[1]) == 1 then return 0 end', + 'redis.call(\'SET\', KEYS[1], ARGV[1], \"PX\", ARGV[3])', + 'redis.call(\'SET\', KEYS[2], ARGV[2], \"PX\", ARGV[3])', + 'return 1', + ].join('\n'); + const result = await this.redis.eval( + script, + 2, + lockKey(workerId), + lockIncarnationKey(workerId), + assignmentId, + incarnationId, + String(ttlSeconds * 1000), + ); + return Number(result) === 1; + } + + private async cleanupDispatch( + workerId: string, + assignmentId: string, + assignment: StoredAssignment | undefined, + ): Promise { + await this.cancel(assignmentId); + if (assignment == null) { + await this.releaseLock(workerId, assignmentId); + return; + } + await this.cleanup(assignment); + } + + private async cleanupWithRetry( + workerId: string, + assignmentId: string, + assignment: StoredAssignment | undefined, + ): Promise { + let lastError: unknown; + for (let attempt = 0; attempt < 3; attempt += 1) { + try { + await this.cleanupDispatch(workerId, assignmentId, assignment); + return; + } catch (error) { + lastError = error; + await delay(25); + } + } + throw lastError; + } + + private async cleanup(assignment: StoredAssignment): Promise { + await Promise.all([ + this.redis.del( + assignmentKey(assignment.assignmentId), + settlementKey(assignment.assignmentId), + ), + this.releaseLock(assignment.workerId, assignment.assignmentId), + ]); + } + + private async releaseLock( + workerId: string, + assignmentId: string, + ): Promise { + const script = [ + 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1], KEYS[2])', + 'end', + 'return 0', + ].join('\n'); + await this.redis.eval( + script, + 2, + lockKey(workerId), + lockIncarnationKey(workerId), + assignmentId, + ); + } +} diff --git a/service/src/config.test.ts b/service/src/config.test.ts index 2f88a87a..22878ea7 100644 --- a/service/src/config.test.ts +++ b/service/src/config.test.ts @@ -14,6 +14,7 @@ describe('sandbox execution configuration', () => { test('accepts every supported backend and session mode', () => { expect(resolveSandboxBackend('http')).toBe('http'); expect(resolveSandboxBackend('lambda-microvm')).toBe('lambda-microvm'); + expect(resolveSandboxBackend('remote-bridge')).toBe('remote-bridge'); expect(resolveRuntimeSessionMode('stateless')).toBe('stateless'); expect(resolveRuntimeSessionMode('affinity')).toBe('affinity'); expect(resolveRuntimeSessionMode('strict')).toBe('strict'); @@ -21,7 +22,7 @@ describe('sandbox execution configuration', () => { test('rejects unknown values instead of silently changing execution semantics', () => { expect(() => resolveSandboxBackend('lambda_microvm')).toThrow( - 'CODEAPI_SANDBOX_BACKEND must be one of: http, lambda-microvm', + 'CODEAPI_SANDBOX_BACKEND must be one of: http, lambda-microvm, remote-bridge', ); expect(() => resolveSandboxBackend('')).toThrow('CODEAPI_SANDBOX_BACKEND'); expect(() => resolveSandboxBackend(' ')).toThrow('CODEAPI_SANDBOX_BACKEND'); diff --git a/service/src/config.ts b/service/src/config.ts index ecf35661..dbc6dfe8 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -243,12 +243,12 @@ function configuredChoice( export function resolveSandboxBackend( raw: string | undefined, -): 'http' | 'lambda-microvm' { +): 'http' | 'lambda-microvm' | 'remote-bridge' { return configuredChoice( raw, 'CODEAPI_SANDBOX_BACKEND', 'http', - ['http', 'lambda-microvm'], + ['http', 'lambda-microvm', 'remote-bridge'], ); } @@ -350,8 +350,13 @@ export const env = { * - `http` (default): POST signed execute requests to SANDBOX_ENDPOINT * (current Kubernetes/libkrun sandbox-runner). * - `lambda-microvm`: AWS Lambda MicroVM backend. + * - `remote-bridge`: dispatch to an outbound-connected @librechat/code worker. */ SANDBOX_BACKEND: sandboxBackend, + /** Outbound worker selected by the remote-bridge backend. */ + BRIDGE_WORKER_ID: process.env.CODEAPI_BRIDGE_WORKER_ID ?? '', + /** Enrollment and lease credential shared only with the configured worker. */ + BRIDGE_TOKEN: process.env.CODEAPI_BRIDGE_TOKEN ?? '', /** * Runtime session affinity for stateful sandbox backends. * - `stateless` (default): no runtime sessions; `runtime_session_hint` ignored. diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 8fc0adc2..c5687610 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -4,6 +4,7 @@ import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from import { validateStartupAuthConfig } from './auth/startup'; import { env } from './config'; import { + validateApiBridgePolicy, validateApiHardenedConfig, validateExecutionProfilePolicy, validateSandboxBackendPolicy, @@ -89,9 +90,12 @@ function setupQueueListeners(queue: Queue, name: string): void { export async function startupApiOnly(): Promise { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); + validateApiBridgePolicy(); validateExecutionProfilePolicy({ requireBackendMatch: false }); /* No validateSandboxBackendPolicy() here: an API-only pod authenticates and * enqueues jobs, it never constructs the Lambda backend or checkpoint store. + * Bridge credentials are validated separately above because this process + * exposes the public registration, lease, and settlement routes. * 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. */ diff --git a/service/src/runtime-session/job-policy.test.ts b/service/src/runtime-session/job-policy.test.ts index 96725f9f..22615f7b 100644 --- a/service/src/runtime-session/job-policy.test.ts +++ b/service/src/runtime-session/job-policy.test.ts @@ -92,6 +92,19 @@ describe('resolveRuntimeSessionForJob', () => { })).toThrow('http/affinity worker cannot honor queued affinity runtime session'); }); + test('allows a remote bridge worker to honor a stateful job', () => { + expect(resolveRuntimeSessionForJob({ + workerBackend: 'remote-bridge', + workerMode: 'strict', + runtimeSessionMode: 'strict', + runtimeSessionId: 'rt_attached', + isSynthetic: false, + })).toEqual({ + runtimeSessionId: 'rt_attached', + runtimeSessionMode: 'strict', + }); + }); + test('rejects contradictory or invalid producer decisions', () => { expect(() => resolveRuntimeSessionForJob({ ...LAMBDA_WORKER, diff --git a/service/src/runtime-session/job-policy.ts b/service/src/runtime-session/job-policy.ts index bdadeddb..8c473c60 100644 --- a/service/src/runtime-session/job-policy.ts +++ b/service/src/runtime-session/job-policy.ts @@ -10,7 +10,7 @@ export type RuntimeSessionJobDecision = { runtimeSessionMode: RuntimeSessionMode; }; -type SandboxBackendName = 'http' | 'lambda-microvm'; +type SandboxBackendName = 'http' | 'lambda-microvm' | 'remote-bridge'; function isRuntimeSessionMode(value: unknown): value is RuntimeSessionMode { return value === 'stateless' || value === 'affinity' || value === 'strict'; @@ -70,7 +70,10 @@ export function resolveRuntimeSessionForJob(args: { if (runtimeSessionId === undefined) { throw new Error(`${runtimeSessionMode} queued job requires a runtimeSessionId`); } - if (args.workerMode === 'stateless' || args.workerBackend !== 'lambda-microvm') { + if ( + args.workerMode === 'stateless' + || (args.workerBackend !== 'lambda-microvm' && args.workerBackend !== 'remote-bridge') + ) { throw new Error( `${args.workerBackend}/${args.workerMode} worker cannot honor queued ` + `${runtimeSessionMode} runtime session`, diff --git a/service/src/sandbox-backend/index.test.ts b/service/src/sandbox-backend/index.test.ts index e9f2ac7d..44378011 100644 --- a/service/src/sandbox-backend/index.test.ts +++ b/service/src/sandbox-backend/index.test.ts @@ -26,6 +26,12 @@ describe('getSandboxBackend', () => { expect(backend.name).toBe('lambda-microvm'); }); + test('selects the outbound remote bridge backend when configured', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + const backend = getSandboxBackend(); + expect(backend.name).toBe('remote-bridge'); + }); + test('does not load Lambda-only modules for the HTTP backend', async () => { const serviceRoot = path.resolve(import.meta.dir, '../..'); const probe = Bun.spawn([ diff --git a/service/src/sandbox-backend/index.ts b/service/src/sandbox-backend/index.ts index 3ebf796a..e5192513 100644 --- a/service/src/sandbox-backend/index.ts +++ b/service/src/sandbox-backend/index.ts @@ -14,6 +14,25 @@ export { HttpSandboxBackend } from './http'; let backend: SandboxBackend | undefined; +class LazyRemoteBridgeSandboxBackend implements SandboxBackend { + readonly name = 'remote-bridge' as const; + private backendPromise: Promise | undefined; + + private load(): Promise { + this.backendPromise ??= import('./remote-bridge').then( + ({ RemoteBridgeSandboxBackend }) => new RemoteBridgeSandboxBackend(), + ); + return this.backendPromise; + } + + async execute( + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + ): Promise { + return (await this.load()).execute(req, ctx); + } +} + class LazyLambdaMicrovmSandboxBackend implements SandboxBackend { readonly name = 'lambda-microvm' as const; private backendPromise: Promise | undefined; @@ -72,6 +91,9 @@ class LazyLambdaMicrovmSandboxBackend implements SandboxBackend { } function createBackend(): SandboxBackend { + if (env.SANDBOX_BACKEND === 'remote-bridge') { + return new LazyRemoteBridgeSandboxBackend(); + } if (env.SANDBOX_BACKEND === 'lambda-microvm') { /* Loading the concrete backend also loads its session registry and * checkpoint code. Defer the whole graph so the default HTTP worker does diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts new file mode 100644 index 00000000..a18788a8 --- /dev/null +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -0,0 +1,82 @@ +import type { + SandboxBackend, + SandboxExecuteContext, + SandboxRawResponse, + SandboxTransportRequest, +} from './types'; +import type { RedisBridgeStore } from '../bridge/store'; + +import { env } from '../config'; +import { bridgeStore } from '../bridge/router'; +import { BridgeStoreError } from '../bridge/store'; +import { SandboxBackendError } from './types'; + +export class RemoteBridgeSandboxBackend implements SandboxBackend { + readonly name = 'remote-bridge' as const; + + constructor( + private readonly store: RedisBridgeStore = bridgeStore, + private readonly workerId: string = env.BRIDGE_WORKER_ID, + ) {} + + async execute( + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + ): Promise { + if (!this.workerId) { + throw new SandboxBackendError( + 'BRIDGE_WORKER_OFFLINE', + 'No bridge worker is configured', + ); + } + const sessionResultFinalizer = ctx.sessionResultFinalizer; + try { + const settlement = await this.store.dispatch({ + workerId: this.workerId, + body: req.body, + headers: req.headers, + runtimeSessionId: ctx.runtimeSessionId, + deadlineAtMs: ctx.deadlineAtMs ?? Date.now() + env.JOB_TIMEOUT, + signal: ctx.signal, + finalize: sessionResultFinalizer + ? async (settlement) => { + if (settlement.status === 'rejected') return settlement; + return { + ...settlement, + result: await sessionResultFinalizer(settlement.result), + }; + } + : undefined, + }); + if (settlement.status === 'rejected') { + throw new SandboxBackendError( + 'BRIDGE_EXECUTION_FAILED', + settlement.error, + ); + } + return settlement.result as SandboxRawResponse; + } catch (error) { + if (!(error instanceof BridgeStoreError)) throw error; + if (error.code === 'WORKER_BUSY') { + throw new SandboxBackendError( + 'BRIDGE_WORKER_BUSY', + error.message, + error, + ); + } + if (error.code === 'ASSIGNMENT_EXPIRED') { + throw new SandboxBackendError( + 'BRIDGE_DEADLINE_EXCEEDED', + error.message, + error, + ); + } + throw new SandboxBackendError( + 'BRIDGE_WORKER_OFFLINE', + error.message, + error, + true, + ); + } + } +} diff --git a/service/src/sandbox-backend/types.ts b/service/src/sandbox-backend/types.ts index 75d7f1af..15bb6942 100644 --- a/service/src/sandbox-backend/types.ts +++ b/service/src/sandbox-backend/types.ts @@ -56,13 +56,17 @@ export type SandboxRawResponse = t.ExecuteResponse & { }; export interface SandboxBackend { - readonly name: 'http' | 'lambda-microvm'; + readonly name: 'http' | 'lambda-microvm' | 'remote-bridge'; execute(req: SandboxTransportRequest, ctx: SandboxExecuteContext): Promise; shutdown?(): Promise; } export type SandboxBackendErrorCode = | 'RUNTIME_SESSION_BUSY' + | 'BRIDGE_WORKER_OFFLINE' + | 'BRIDGE_WORKER_BUSY' + | 'BRIDGE_EXECUTION_FAILED' + | 'BRIDGE_DEADLINE_EXCEEDED' | 'MICROVM_LAUNCH_FAILED' | 'MICROVM_LAUNCH_THROTTLED' | 'MICROVM_UNHEALTHY' diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 4aa603e4..56e0fb23 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -1,6 +1,7 @@ import { afterEach, describe, expect, test } from 'bun:test'; import { env } from './config'; import { + validateApiBridgePolicy, validateApiHardenedConfig, validateEgressGatewayHardenedConfig, validateExecutionProfilePolicy, @@ -14,6 +15,8 @@ const saved = { executionProfile: env.EXECUTION_PROFILE, executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: env.SANDBOX_BACKEND, + bridgeWorkerId: env.BRIDGE_WORKER_ID, + bridgeToken: env.BRIDGE_TOKEN, ptcMode: env.PTC_MODE, runtimeSessionMode: env.RUNTIME_SESSION_MODE, lambdaImageArn: env.LAMBDA_MICROVM_IMAGE_ARN, @@ -52,6 +55,8 @@ function restore(): void { env.EXECUTION_PROFILE = saved.executionProfile; env.EXECUTION_PROFILE_SOURCE = saved.executionProfileSource; env.SANDBOX_BACKEND = saved.sandboxBackend; + env.BRIDGE_WORKER_ID = saved.bridgeWorkerId; + env.BRIDGE_TOKEN = saved.bridgeToken; env.PTC_MODE = saved.ptcMode; env.RUNTIME_SESSION_MODE = saved.runtimeSessionMode; env.LAMBDA_MICROVM_IMAGE_ARN = saved.lambdaImageArn; @@ -279,12 +284,62 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); - test('stateful runtime session modes require the lambda backend', () => { + test('stateful runtime session modes require a stateful backend', () => { env.SANDBOX_BACKEND = 'http'; env.RUNTIME_SESSION_MODE = 'affinity'; - expect(() => validateSandboxBackendPolicy()).toThrow('requires the lambda-microvm backend'); + expect(() => validateSandboxBackendPolicy()).toThrow( + 'requires the lambda-microvm or remote-bridge backend', + ); + env.RUNTIME_SESSION_MODE = 'strict'; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'requires the lambda-microvm or remote-bridge backend', + ); + }); + + test('accepts a configured remote bridge and fails closed on missing enrollment', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; env.RUNTIME_SESSION_MODE = 'strict'; - expect(() => validateSandboxBackendPolicy()).toThrow('requires the lambda-microvm backend'); + env.PTC_MODE = 'replay'; + env.BRIDGE_WORKER_ID = ''; + env.BRIDGE_TOKEN = ''; + expect(() => validateSandboxBackendPolicy()).toThrow('CODEAPI_BRIDGE_WORKER_ID'); + + env.BRIDGE_WORKER_ID = 'engineering-vm'; + expect(() => validateSandboxBackendPolicy()).toThrow('CODEAPI_BRIDGE_TOKEN'); + + env.BRIDGE_TOKEN = 'development-bridge-token'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('remote bridge requires replay PTC and a strong token in hardened mode', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.RUNTIME_SESSION_MODE = 'affinity'; + env.BRIDGE_WORKER_ID = 'engineering-vm'; + env.BRIDGE_TOKEN = 'development-bridge-token'; + env.PTC_MODE = 'blocking'; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'PTC replay is the only supported PTC mode', + ); + + env.PTC_MODE = 'replay'; + env.HARDENED_SANDBOX_MODE = true; + expect(() => validateSandboxBackendPolicy()).toThrow('at least 32 bytes'); + + env.BRIDGE_TOKEN = 'strong-remote-bridge-token-32-bytes'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('API bridge policy requires a strong token in hardened mode', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.BRIDGE_WORKER_ID = 'engineering-vm'; + env.BRIDGE_TOKEN = 'short-token'; + env.PTC_MODE = 'replay'; + env.HARDENED_SANDBOX_MODE = true; + + expect(() => validateApiBridgePolicy()).toThrow('at least 32 bytes'); + + env.BRIDGE_TOKEN = 'strong-remote-bridge-token-32-bytes'; + expect(() => validateApiBridgePolicy()).not.toThrow(); }); test('rejects blocking PTC on the lambda backend', () => { diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index a82dabf0..c324cde4 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -53,6 +53,22 @@ export function validateApiHardenedConfig(): void { requireValue(INTERNAL_SERVICE_TOKEN_ENV, process.env[INTERNAL_SERVICE_TOKEN_ENV]); } +/** Validate bridge credentials in every process that exposes bridge routes. */ +export function validateApiBridgePolicy(): void { + if (env.SANDBOX_BACKEND !== 'remote-bridge') return; + requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); + if (env.HARDENED_SANDBOX_MODE) { + requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); + } else { + requireValue('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); + } + if (env.PTC_MODE === 'blocking') { + throw new SecureStartupConfigError( + 'PTC replay is the only supported PTC mode for the remote-bridge backend (unset PTC_MODE=blocking)', + ); + } +} + export function validateWorkerHardenedConfig(): void { if (!env.HARDENED_SANDBOX_MODE) return; rejectValue('CODEAPI_EGRESS_GRANT_SECRET', process.env.CODEAPI_EGRESS_GRANT_SECRET); @@ -93,11 +109,15 @@ export function validateExecutionProfilePolicy(options: { if ( env.RUNTIME_SESSION_MODE === 'stateless' - || (requireBackendMatch && env.SANDBOX_BACKEND !== 'lambda-microvm') + || ( + requireBackendMatch + && env.SANDBOX_BACKEND !== 'lambda-microvm' + && env.SANDBOX_BACKEND !== 'remote-bridge' + ) ) { throw new SecureStartupConfigError( 'CODEAPI_EXECUTION_PROFILE=stateful requires ' - + (requireBackendMatch ? 'CODEAPI_SANDBOX_BACKEND=lambda-microvm and ' : '') + + (requireBackendMatch ? 'CODEAPI_SANDBOX_BACKEND=lambda-microvm or remote-bridge and ' : '') + 'CODEAPI_RUNTIME_SESSION_MODE=affinity or strict', ); } @@ -111,9 +131,13 @@ export function validateSandboxBackendPolicy(): void { if (env.RUNTIME_SESSION_MODE !== 'stateless' && env.SANDBOX_BACKEND === 'http') { throw new SecureStartupConfigError( `CODEAPI_RUNTIME_SESSION_MODE=${env.RUNTIME_SESSION_MODE} requires ` - + 'the lambda-microvm backend; use stateless mode with the http backend', + + 'the lambda-microvm or remote-bridge backend; use stateless mode with the http backend', ); } + if (env.SANDBOX_BACKEND === 'remote-bridge') { + validateApiBridgePolicy(); + return; + } if (env.SANDBOX_BACKEND !== 'lambda-microvm') return; const numericConfigError = lambdaMicrovmNumericConfigError(env); diff --git a/service/src/service-api.ts b/service/src/service-api.ts index ec15b1ac..6a0a8331 100644 --- a/service/src/service-api.ts +++ b/service/src/service-api.ts @@ -5,6 +5,7 @@ import { requestErrorLogger, requestNotFoundLogger } from './middleware/request- import { executionProfileMiddleware } from './middleware/execution-profile'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; +import bridgeRouter from './bridge/router'; import { connection } from './queue'; import { env } from './config'; import logger from './logger'; @@ -28,6 +29,7 @@ app.get('/v1/health', async (_, res) => { } }); +v1.use('/bridge', bridgeRouter); v1.use(apiKeyAuth); v1.use(serviceRouter); diff --git a/service/src/utils.test.ts b/service/src/utils.test.ts index f1952dfc..8aa0dac0 100644 --- a/service/src/utils.test.ts +++ b/service/src/utils.test.ts @@ -145,6 +145,30 @@ describe('sandbox error formatting', () => { }); }); + test('maps remote bridge failures without exposing worker details', () => { + const cases = [ + ['BRIDGE_WORKER_OFFLINE', 503, 'Remote code worker is unavailable'], + ['BRIDGE_WORKER_BUSY', 409, 'Remote code worker is busy'], + ['BRIDGE_EXECUTION_FAILED', 502, 'Remote code execution failed'], + [ + 'BRIDGE_DEADLINE_EXCEEDED', + 504, + 'Remote code execution deadline exceeded', + ], + ] as const; + for (const [code, status, message] of cases) { + const failure = publicExecutionFailure( + new Error(`${code}: worker vm-private failed at redis.internal`), + ); + expect(failure).toEqual({ + status, + body: { error: code.toLowerCase(), message }, + }); + expect(JSON.stringify(failure)).not.toContain('vm-private'); + expect(JSON.stringify(failure)).not.toContain('redis.internal'); + } + }); + test('maps a recycled dirty session to a retryable public failure', () => { const failure = publicExecutionFailure( new Error('MICROVM_UNHEALTHY: Runtime session rt_private workspace was dirty and has been recycled'), diff --git a/service/src/utils.ts b/service/src/utils.ts index 3078d6e8..b94d8fd2 100644 --- a/service/src/utils.ts +++ b/service/src/utils.ts @@ -128,15 +128,19 @@ export function publicExecutionFailure(error: unknown): { status: number; body: } /* Typed worker failures cross BullMQ as `: `. Runtime-session - * and MicroVM codes describe sandbox availability; SESSION_INPUT_* codes + * MicroVM, and bridge codes describe sandbox availability; SESSION_INPUT_* codes * describe the caller's declared input set or its upstream object source. */ const backendMatch = message.match( - /^(RUNTIME_SESSION_BUSY|MICROVM_[A-Z_]+|SESSION_INPUT_[A-Z_]+):\s*(.+)$/, + /^(RUNTIME_SESSION_BUSY|MICROVM_[A-Z_]+|BRIDGE_[A-Z_]+|SESSION_INPUT_[A-Z_]+):\s*(.+)$/, ); if (backendMatch) { const code = backendMatch[1]; const statuses: Record = { RUNTIME_SESSION_BUSY: 409, + BRIDGE_WORKER_OFFLINE: 503, + BRIDGE_WORKER_BUSY: 409, + BRIDGE_EXECUTION_FAILED: 502, + BRIDGE_DEADLINE_EXCEEDED: 504, SESSION_INPUT_TOO_LARGE: 413, SESSION_INPUT_UNAVAILABLE: 422, SESSION_INPUT_SOURCE_FAILED: 502, @@ -147,6 +151,10 @@ export function publicExecutionFailure(error: unknown): { status: number; body: const status = statuses[code] ?? (sessionInputFailure ? 500 : 503); const publicMessages: Record = { RUNTIME_SESSION_BUSY: 'Runtime session is busy', + BRIDGE_WORKER_OFFLINE: 'Remote code worker is unavailable', + BRIDGE_WORKER_BUSY: 'Remote code worker is busy', + BRIDGE_EXECUTION_FAILED: 'Remote code execution failed', + BRIDGE_DEADLINE_EXCEEDED: 'Remote code execution deadline exceeded', MICROVM_LAUNCH_FAILED: 'Sandbox launch failed', MICROVM_LAUNCH_THROTTLED: 'Sandbox capacity is temporarily unavailable', MICROVM_UNHEALTHY: 'Sandbox runtime is unavailable', diff --git a/service/tsconfig.json b/service/tsconfig.json index c3e88635..dcddf82d 100644 --- a/service/tsconfig.json +++ b/service/tsconfig.json @@ -13,7 +13,11 @@ "@/*": ["src/*"] } }, - "include": ["src/**/*.ts", "../shared/telemetry-core.ts"], + "include": [ + "src/**/*.ts", + "../shared/telemetry-core.ts", + "../packages/code/src/protocol.ts" + ], "exclude": [ "node_modules", "**/*.spec.ts",