From 7860432315ae60ebc4c619851af4b0c6c6df2004 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 01:38:23 -0400 Subject: [PATCH 01/29] feat: add outbound stateful code bridge --- .env.example | 7 + .gitignore | 1 + README.md | 14 + api/Dockerfile | 5 +- api/src/entrypoint.sh | 9 +- docker/Dockerfile.worker-sandbox | 5 +- docker/rootfs-setup.c | 92 +++++ docker/start-direct-sandbox.sh | 33 +- docs/remote-bridge/README.md | 83 +++++ packages/code/Dockerfile | 15 + packages/code/README.md | 33 ++ packages/code/package-lock.json | 54 +++ packages/code/package.json | 42 +++ packages/code/src/cli.ts | 47 +++ packages/code/src/index.ts | 2 + packages/code/src/protocol.test.ts | 10 + packages/code/src/protocol.ts | 87 +++++ packages/code/src/worker.test.ts | 70 ++++ packages/code/src/worker.ts | 225 ++++++++++++ packages/code/tsconfig.json | 15 + service/Dockerfile | 2 + service/Dockerfile.local | 2 + service/rollup.config.js | 6 +- service/src/api-server.ts | 2 + service/src/bridge/router.ts | 186 ++++++++++ service/src/bridge/store.test.ts | 105 ++++++ service/src/bridge/store.ts | 325 ++++++++++++++++++ service/src/config.test.ts | 3 +- service/src/config.ts | 9 +- .../src/runtime-session/job-policy.test.ts | 13 + service/src/runtime-session/job-policy.ts | 7 +- service/src/sandbox-backend/index.test.ts | 6 + service/src/sandbox-backend/index.ts | 22 ++ service/src/sandbox-backend/remote-bridge.ts | 72 ++++ service/src/sandbox-backend/types.ts | 6 +- service/src/secure-startup.ts | 24 +- service/src/service-api.ts | 2 + service/tsconfig.json | 6 +- 38 files changed, 1600 insertions(+), 47 deletions(-) create mode 100644 docker/rootfs-setup.c create mode 100644 docs/remote-bridge/README.md create mode 100644 packages/code/Dockerfile create mode 100644 packages/code/README.md create mode 100644 packages/code/package-lock.json create mode 100644 packages/code/package.json create mode 100644 packages/code/src/cli.ts create mode 100644 packages/code/src/index.ts create mode 100644 packages/code/src/protocol.test.ts create mode 100644 packages/code/src/protocol.ts create mode 100644 packages/code/src/worker.test.ts create mode 100644 packages/code/src/worker.ts create mode 100644 packages/code/tsconfig.json create mode 100644 service/src/bridge/router.ts create mode 100644 service/src/bridge/store.test.ts create mode 100644 service/src/bridge/store.ts create mode 100644 service/src/sandbox-backend/remote-bridge.ts 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..5334e360 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 @@ -231,6 +233,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..14161547 --- /dev/null +++ b/docs/remote-bridge/README.md @@ -0,0 +1,83 @@ +# 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). + +## 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. +- 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/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..938304a4 --- /dev/null +++ b/packages/code/README.md @@ -0,0 +1,33 @@ +# `@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`. + +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..23005cb3 --- /dev/null +++ b/packages/code/src/cli.ts @@ -0,0 +1,47 @@ +#!/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 worker = new BridgeWorker({ + codeApiUrl: required('LIBRECHAT_CODE_URL'), + token: required('LIBRECHAT_CODE_WORKER_TOKEN'), + workerId: required('LIBRECHAT_CODE_WORKER_ID'), + sandboxEndpoint: + process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? + 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + 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..7cb757a2 --- /dev/null +++ b/packages/code/src/protocol.ts @@ -0,0 +1,87 @@ +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; + capabilities: BridgeWorkerCapabilities; +} + +export interface BridgeWorkerRegistrationResponse { + protocolVersion: BridgeProtocolVersion; + workerId: string; + registeredAt: string; + leaseTtlMs: number; +} + +export interface BridgeSandboxRequest { + body: TBody; + headers: Record; +} + +export interface BridgeAssignment { + protocolVersion: BridgeProtocolVersion; + assignmentId: string; + workerId: 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; + status: 'fulfilled'; + result: TResult; +} + +export interface BridgeRejectedSettlement { + protocolVersion: BridgeProtocolVersion; + generation: number; + leaseToken: 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..0e89babf --- /dev/null +++ b/packages/code/src/worker.test.ts @@ -0,0 +1,70 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import { BridgeWorker } 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', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2/', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + const assignment: BridgeAssignment = { + protocolVersion: 1, + assignmentId: 'assignment-1', + workerId: 'vm-1', + 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/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', + status: 'fulfilled', + result: { session_id: 'run-1', files: [] }, + }); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts new file mode 100644 index 00000000..b5767418 --- /dev/null +++ b/packages/code/src/worker.ts @@ -0,0 +1,225 @@ +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; +} + +const DEFAULT_LEASE_WAIT_MS = 25_000; +const DEFAULT_RECONNECT_DELAY_MS = 1_000; + +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 BridgeWorker { + private readonly fetchImpl: typeof fetch; + private readonly codeApiUrl: string; + private readonly sandboxEndpoint: string; + + constructor(private readonly options: BridgeWorkerOptions) { + this.fetchImpl = options.fetchImpl ?? fetch; + this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); + this.sandboxEndpoint = normalizedBaseUrl(options.sandboxEndpoint); + } + + async register( + signal?: AbortSignal, + ): Promise { + return this.request( + `${this.codeApiUrl}/bridge/workers/register`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: this.options.workerId, + capabilities: this.options.capabilities, + }, + signal, + ); + } + + 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, + }, + signal, + ); + 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 (signal?.aborted) return; + if ( + error instanceof BridgeProtocolError && + (error.status === 401 || error.status === 403) + ) { + 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 cancellationController = new AbortController(); + const cancellationWatcher = this.watchCancellation( + assignment, + executionController, + cancellationController.signal, + ); + let settlement: BridgeSettlement; + try { + const headers = { + ...assignment.request.headers, + ...(assignment.runtimeSessionId + ? { 'X-Runtime-Session-Id': assignment.runtimeSessionId } + : {}), + }; + const response = await this.fetchImpl(`${this.sandboxEndpoint}/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 (!response.ok) { + 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, + status: 'fulfilled', + result: payload, + }; + } catch (error) { + settlement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected', + error: + error instanceof Error ? error.message : 'Sandbox execution failed', + }; + } + + cancellationController.abort(); + await cancellationWatcher; + signal?.removeEventListener('abort', abortExecution); + await this.request( + this.assignmentUrl(assignment, 'settle'), + settlement, + signal, + ); + } + + private assignmentUrl(assignment: BridgeAssignment, action: string): string { + return ( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + + `/assignments/${encodeURIComponent(assignment.assignmentId)}/${action}` + ); + } + + 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 }, + 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.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/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..0e328fb7 --- /dev/null +++ b/service/src/bridge/router.ts @@ -0,0 +1,186 @@ +import { timingSafeEqual } from 'crypto'; + +import { Router } from 'express'; +import type { NextFunction, Request, 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 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 isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function sendStoreError(error: BridgeStoreError, res: Response): void { + const status = error.code === 'ASSIGNMENT_NOT_FOUND' ? 404 : 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 + ) { + 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', async (req: Request, res: Response) => { + const registration = req.body as unknown; + if ( + !isRecord(registration) || + registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof registration.workerId !== 'string' || + !validWorkerId(registration.workerId) || + !isRecord(registration.capabilities) || + registration.capabilities.statefulWorkspace !== true || + 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; + } + await bridgeStore.register(registration as unknown as BridgeWorkerRegistration); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); +}); + +router.post('/workers/:workerId/lease', async (req: Request, res: Response) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + const requestedWait = Number(body.waitMs ?? 25_000); + if ( + !validWorkerId(workerId) || + !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; + } + const assignment = await bridgeStore.lease( + workerId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); +}); + +router.post( + '/workers/:workerId/assignments/:assignmentId/settle', + 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', + async (req, res) => { + const cancelled = await bridgeStore.cancelled( + req.params.workerId, + 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..2f6a12ee --- /dev/null +++ b/service/src/bridge/store.test.ts @@ -0,0 +1,105 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +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); + +afterEach(async () => { + await redis.flushall(); +}); + +describe('RedisBridgeStore', () => { + test('delivers and settles one fenced stateful assignment', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + 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', 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 ?? '', + 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', + 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', 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', + status: 'rejected', + error: 'unused', + }), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); +}); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts new file mode 100644 index 00000000..d6f98aa2 --- /dev/null +++ b/service/src/bridge/store.ts @@ -0,0 +1,325 @@ +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; +const MAX_ASSIGNMENT_TTL_SECONDS = 10 * 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_MISMATCH', + message: string, + ) { + super(message); + this.name = 'BridgeStoreError'; + } +} + +interface StoredAssignment extends CodeBridgeAssignment { + leaseTokenHash: string; +} + +function workerKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}`; +} + +function queueKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:assignments`; +} + +function generationKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:generation`; +} + +function lockKey(workerId: string): string { + return `${PREFIX}:worker:${workerId}:lock`; +} + +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.min( + MAX_ASSIGNMENT_TTL_SECONDS, + 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 timer = setTimeout(resolve, ms); + signal?.addEventListener( + 'abort', + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + }); +} + +export class RedisBridgeStore { + constructor( + private readonly redis: Redis, + private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, + ) {} + + async register(registration: BridgeWorkerRegistration): Promise { + await this.redis.set( + workerKey(registration.workerId), + JSON.stringify(registration), + 'EX', + this.workerTtlSeconds, + ); + } + + async dispatch(args: { + workerId: string; + body: t.PayloadBody; + headers: Record; + runtimeSessionId?: string; + deadlineAtMs: number; + signal: AbortSignal; + }): Promise { + const 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`, + ); + } + + const assignmentId = randomBytes(18).toString('base64url'); + const leaseToken = randomBytes(32).toString('base64url'); + const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs); + const locked = await this.redis.set( + lockKey(args.workerId), + assignmentId, + 'PX', + ttlSeconds * 1000, + 'NX', + ); + if (locked !== 'OK') { + throw new BridgeStoreError( + 'WORKER_BUSY', + `Bridge worker ${args.workerId} is busy`, + ); + } + + const generation = await this.redis.incr(generationKey(args.workerId)); + const assignment: StoredAssignment = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + assignmentId, + workerId: args.workerId, + generation, + leaseToken, + leaseTokenHash: tokenHash(leaseToken), + expiresAt: new Date(args.deadlineAtMs).toISOString(), + runtimeSessionId: args.runtimeSessionId, + request: { + body: args.body, + headers: args.headers, + }, + }; + + try { + const transaction = this.redis.multi(); + transaction.set( + assignmentKey(assignmentId), + JSON.stringify(assignment), + 'EX', + ttlSeconds, + ); + transaction.rpush(queueKey(args.workerId), assignmentId); + transaction.expire(queueKey(args.workerId), ttlSeconds); + await transaction.exec(); + return await this.waitForSettlement( + assignment, + args.deadlineAtMs, + args.signal, + ); + } finally { + await this.cancel(assignmentId); + await this.cleanup(assignment); + } + } + + async lease( + workerId: 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)); + 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 (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', + ); + } + if ( + 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)); + await this.redis.set( + settlementKey(assignmentId), + JSON.stringify(settlement), + 'EX', + ttlSeconds, + ); + } + + async cancelled(workerId: string, assignmentId: string): Promise { + const assignment = await this.readAssignment(assignmentId); + if (assignment == null || assignment.workerId !== workerId) return true; + return (await this.redis.exists(cancellationKey(assignmentId))) === 1; + } + + 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); + } + 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 cleanup(assignment: StoredAssignment): Promise { + const script = [ + 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1])', + 'end', + 'return 0', + ].join('\n'); + await Promise.all([ + this.redis.del( + assignmentKey(assignment.assignmentId), + settlementKey(assignment.assignmentId), + ), + this.redis.eval( + script, + 1, + lockKey(assignment.workerId), + assignment.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/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..671d1850 --- /dev/null +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -0,0 +1,72 @@ +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', + ); + } + 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, + }); + 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.ts b/service/src/secure-startup.ts index a82dabf0..dc9cb4ca 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -93,11 +93,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 +115,23 @@ 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') { + 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)', + ); + } + 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/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", From 7cf93507726119148d21ca6269e5431e05547abd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 01:45:43 -0400 Subject: [PATCH 02/29] test: cover remote bridge startup policy --- service/src/secure-startup.test.ts | 47 ++++++++++++++++++++++++++++-- 1 file changed, 44 insertions(+), 3 deletions(-) diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 4aa603e4..113b4fc9 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -14,6 +14,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 +54,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 +283,49 @@ 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 backend'); + 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'; + 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('rejects blocking PTC on the lambda backend', () => { From 9be996f4ae1e33aeb5be87158e2ff31a888a0790 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:04:27 -0400 Subject: [PATCH 03/29] fix: harden remote bridge lifecycle fencing --- docker/Dockerfile.worker-sandbox | 1 + docs/remote-bridge/README.md | 4 + launcher/Dockerfile | 7 +- packages/code/README.md | 10 + packages/code/src/cli.ts | 16 +- packages/code/src/protocol.ts | 5 + packages/code/src/worker.test.ts | 115 +++++++++- packages/code/src/worker.ts | 126 ++++++++++- service/Dockerfile.api | 2 + service/Dockerfile.worker | 2 + service/src/bridge/router.ts | 54 ++++- service/src/bridge/store.test.ts | 176 +++++++++++++++- service/src/bridge/store.ts | 209 +++++++++++++++---- service/src/sandbox-backend/remote-bridge.ts | 10 + service/src/utils.test.ts | 24 +++ service/src/utils.ts | 12 +- 16 files changed, 705 insertions(+), 68 deletions(-) diff --git a/docker/Dockerfile.worker-sandbox b/docker/Dockerfile.worker-sandbox index 5334e360..c18eab82 100644 --- a/docker/Dockerfile.worker-sandbox +++ b/docker/Dockerfile.worker-sandbox @@ -85,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/*' diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 14161547..f5632d5d 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -33,6 +33,10 @@ 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 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/README.md b/packages/code/README.md index 938304a4..a79a1709 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -27,6 +27,16 @@ Optional environment variables: - `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. + +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. 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 diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 23005cb3..411dda7f 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -22,15 +22,23 @@ 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: - process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? - 'http://127.0.0.1:2000/api/v2', + sandboxEndpoint, capabilities: { - statefulWorkspace: true, + statefulWorkspace, sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? 'nsjail', runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), policyDigest: createHash('sha256').update(policy).digest('hex'), diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 7cb757a2..a79d4ff4 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -12,12 +12,14 @@ export interface BridgeWorkerCapabilities { export interface BridgeWorkerRegistration { protocolVersion: BridgeProtocolVersion; workerId: string; + incarnationId: string; capabilities: BridgeWorkerCapabilities; } export interface BridgeWorkerRegistrationResponse { protocolVersion: BridgeProtocolVersion; workerId: string; + incarnationId: string; registeredAt: string; leaseTtlMs: number; } @@ -31,6 +33,7 @@ export interface BridgeAssignment { protocolVersion: BridgeProtocolVersion; assignmentId: string; workerId: string; + incarnationId: string; generation: number; leaseToken: string; expiresAt: string; @@ -47,6 +50,7 @@ export interface BridgeFulfilledSettlement { protocolVersion: BridgeProtocolVersion; generation: number; leaseToken: string; + incarnationId: string; status: 'fulfilled'; result: TResult; } @@ -55,6 +59,7 @@ export interface BridgeRejectedSettlement { protocolVersion: BridgeProtocolVersion; generation: number; leaseToken: string; + incarnationId: string; status: 'rejected'; error: string; } diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 0e89babf..03061d91 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -27,7 +27,9 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' codeApiUrl: 'https://code.example/v1/', token: 'worker-secret', workerId: 'vm-1', - sandboxEndpoint: 'http://127.0.0.1:2000/api/v2/', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2/', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -39,6 +41,7 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' 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(), @@ -52,7 +55,10 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' await worker.executeAndSettle(assignment); assert.equal(requests.length, 2); - assert.equal(requests[0].url, 'http://127.0.0.1:2000/api/v2/execute'); + 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' @@ -64,7 +70,112 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' 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); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index b5767418..31083947 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -1,3 +1,5 @@ +import { randomBytes } from 'node:crypto'; + import { BRIDGE_PROTOCOL_VERSION, BridgeProtocolError, @@ -23,10 +25,14 @@ export interface BridgeWorkerOptions { 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 RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; function normalizedBaseUrl(value: string): string { return value.replace(/\/+$/, ''); @@ -41,25 +47,37 @@ 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 { - return this.request( + 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 { @@ -68,9 +86,18 @@ export class BridgeWorker { { 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; } @@ -85,7 +112,7 @@ export class BridgeWorker { if (signal?.aborted) return; if ( error instanceof BridgeProtocolError && - (error.status === 401 || error.status === 403) + (error.status === 401 || error.status === 403 || error.status === 409) ) { throw error; } @@ -104,6 +131,23 @@ export class BridgeWorker { 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, @@ -118,16 +162,20 @@ export class BridgeWorker { ? { 'X-Runtime-Session-Id': assignment.runtimeSessionId } : {}), }; - const response = await this.fetchImpl(`${this.sandboxEndpoint}/execute`, { - method: 'POST', - headers: { - ...headers, - 'Content-Type': 'application/json', + 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, }, - body: JSON.stringify(assignment.request.body), - signal: executionController.signal, - }); + ); const payload = (await response.json()) as object; + if (heartbeatError != null) throw heartbeatError; if (!response.ok) { throw new BridgeProtocolError( errorMessage(payload) ?? @@ -139,6 +187,7 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, status: 'fulfilled', result: payload, }; @@ -147,12 +196,16 @@ export class BridgeWorker { 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); + heartbeatController.abort(); + await heartbeat; cancellationController.abort(); await cancellationWatcher; signal?.removeEventListener('abort', abortExecution); @@ -163,6 +216,54 @@ export class BridgeWorker { ); } + private sandboxEndpointFor(assignment: BridgeAssignment): string { + if (assignment.runtimeSessionId == null) return this.sandboxEndpoint; + 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)}` + @@ -181,7 +282,10 @@ export class BridgeWorker { try { const response = await this.request<{ cancelled: boolean }>( this.assignmentUrl(assignment, 'cancellation'), - { protocolVersion: BRIDGE_PROTOCOL_VERSION }, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + }, signal, ); if (response.cancelled) { 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.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/src/bridge/router.ts b/service/src/bridge/router.ts index 0e328fb7..c69e1e03 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -11,6 +11,7 @@ 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); @@ -45,6 +46,10 @@ 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; } @@ -62,7 +67,8 @@ function isSettlement(value: unknown): value is CodeBridgeSettlement { !Number.isSafeInteger(value.generation) || value.generation < 1 || typeof value.leaseToken !== 'string' || - value.leaseToken.length < 32 + value.leaseToken.length < 32 || + !validIncarnationId(value.incarnationId) ) { return false; } @@ -86,8 +92,9 @@ router.post('/workers/register', async (req: Request, res: Response) => { registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || typeof registration.workerId !== 'string' || !validWorkerId(registration.workerId) || + !validIncarnationId(registration.incarnationId) || !isRecord(registration.capabilities) || - registration.capabilities.statefulWorkspace !== true || + typeof registration.capabilities.statefulWorkspace !== 'boolean' || typeof registration.capabilities.sandboxProfile !== 'string' || registration.capabilities.sandboxProfile.trim().length === 0 || registration.capabilities.sandboxProfile.length > 128 || @@ -110,10 +117,21 @@ router.post('/workers/register', async (req: Request, res: Response) => { }); return; } - await bridgeStore.register(registration as unknown as BridgeWorkerRegistration); + 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, }); @@ -125,6 +143,8 @@ router.post('/workers/:workerId/lease', async (req: Request, res: Response) => { const requestedWait = Number(body.waitMs ?? 25_000); if ( !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || !Number.isFinite(requestedWait) || requestedWait < 0 ) { @@ -137,11 +157,20 @@ router.post('/workers/:workerId/lease', async (req: Request, res: Response) => { }); return; } - const assignment = await bridgeStore.lease( - workerId, - Math.min(requestedWait, MAX_LEASE_WAIT_MS), - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + 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( @@ -175,8 +204,17 @@ router.post( router.post( '/workers/:workerId/assignments/:assignmentId/cancellation', 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 }); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 2f6a12ee..c35a8fd8 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -7,6 +7,7 @@ import { RedisBridgeStore } from './store'; const redis = new RedisMock() as unknown as Redis; const store = new RedisBridgeStore(redis); +const incarnationId = 'incarnation-00000001'; afterEach(async () => { await redis.flushall(); @@ -17,6 +18,7 @@ describe('RedisBridgeStore', () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'vm-1', + incarnationId, capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -32,7 +34,7 @@ describe('RedisBridgeStore', () => { deadlineAtMs: Date.now() + 5_000, signal: controller.signal, }); - const assignment = await store.lease('vm-1', 1_000); + const assignment = await store.lease('vm-1', incarnationId, 1_000); expect(assignment).toBeDefined(); expect(assignment?.runtimeSessionId).toBe('rt-user-1'); @@ -40,6 +42,7 @@ describe('RedisBridgeStore', () => { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment?.generation ?? 0, leaseToken: assignment?.leaseToken ?? '', + incarnationId, status: 'fulfilled', result: { language: 'bash', @@ -72,6 +75,7 @@ describe('RedisBridgeStore', () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'vm-1', + incarnationId, capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -86,13 +90,14 @@ describe('RedisBridgeStore', () => { deadlineAtMs: Date.now() + 5_000, signal: controller.signal, }); - const assignment = await store.lease('vm-1', 1_000); + 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', }), @@ -102,4 +107,171 @@ describe('RedisBridgeStore', () => { 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('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 index d6f98aa2..04fb8e29 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -32,6 +32,9 @@ export class BridgeStoreError extends Error { | 'ASSIGNMENT_EXPIRED' | 'ASSIGNMENT_FENCED' | 'ASSIGNMENT_NOT_FOUND' + | 'WORKER_FENCED' + | 'WORKER_QUARANTINED' + | 'WORKSPACE_QUARANTINED' | 'WORKER_MISMATCH', message: string, ) { @@ -48,6 +51,28 @@ 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): string { return `${PREFIX}:worker:${workerId}:assignments`; } @@ -108,12 +133,45 @@ export class RedisBridgeStore { ) {} async register(registration: BridgeWorkerRegistration): Promise { - await this.redis.set( - workerKey(registration.workerId), - JSON.stringify(registration), - 'EX', - this.workerTtlSeconds, + 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 current then', + ' if current ~= ARGV[1] then', + ' 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, + 4, + workerKey(registration.workerId), + incarnationFenceKey(registration.workerId, registration.incarnationId), + quarantineKey(registration.workerId, registration.incarnationId), + workerIncarnationKey(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', + ); + } } async dispatch(args: { @@ -123,6 +181,9 @@ export class RedisBridgeStore { runtimeSessionId?: string; deadlineAtMs: number; signal: AbortSignal; + finalize?: ( + settlement: CodeBridgeSettlement, + ) => Promise; }): Promise { const registration = await this.registration(args.workerId); if (registration == null) { @@ -140,6 +201,17 @@ export class RedisBridgeStore { `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'); @@ -158,23 +230,24 @@ export class RedisBridgeStore { ); } - const generation = await this.redis.incr(generationKey(args.workerId)); - const assignment: StoredAssignment = { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - assignmentId, - workerId: args.workerId, - generation, - leaseToken, - leaseTokenHash: tokenHash(leaseToken), - expiresAt: new Date(args.deadlineAtMs).toISOString(), - runtimeSessionId: args.runtimeSessionId, - request: { - body: args.body, - headers: args.headers, - }, - }; - + let assignment: StoredAssignment | undefined; 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, + }, + }; const transaction = this.redis.multi(); transaction.set( assignmentKey(assignmentId), @@ -185,19 +258,37 @@ export class RedisBridgeStore { transaction.rpush(queueKey(args.workerId), assignmentId); transaction.expire(queueKey(args.workerId), ttlSeconds); await transaction.exec(); - return await this.waitForSettlement( + const settlement = await this.waitForSettlement( assignment, args.deadlineAtMs, args.signal, ); + if (args.finalize == null) return settlement; + try { + return await args.finalize(settlement); + } catch (error) { + await this.quarantine(args.workerId, registration.incarnationId); + if (args.runtimeSessionId !== undefined) { + await this.redis.set( + workspaceQuarantineKey(args.workerId, args.runtimeSessionId), + '1', + ); + } + throw error; + } } finally { await this.cancel(assignmentId); - await this.cleanup(assignment); + if (assignment == null) { + await this.releaseLock(args.workerId, assignmentId); + } else { + await this.cleanup(assignment); + } } } async lease( workerId: string, + incarnationId: string, waitMs: number, signal?: AbortSignal, ): Promise { @@ -213,6 +304,14 @@ export class RedisBridgeStore { } 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; @@ -238,7 +337,10 @@ export class RedisBridgeStore { '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 ) { @@ -262,12 +364,43 @@ export class RedisBridgeStore { ); } - async cancelled(workerId: string, assignmentId: string): Promise { + async cancelled( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { const assignment = await this.readAssignment(assignmentId); - if (assignment == null || assignment.workerId !== workerId) return true; + 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): Promise { + const script = [ + 'redis.call(\'SET\', KEYS[2], \"1\")', + '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'); + await this.redis.eval( + script, + 3, + workerKey(workerId), + quarantineKey(workerId, incarnationId), + workerIncarnationKey(workerId), + incarnationId, + ); + } + private async registration( workerId: string, ): Promise { @@ -303,23 +436,25 @@ export class RedisBridgeStore { } private async cleanup(assignment: StoredAssignment): Promise { - const script = [ - 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', - ' return redis.call(\'DEL\', KEYS[1])', - 'end', - 'return 0', - ].join('\n'); await Promise.all([ this.redis.del( assignmentKey(assignment.assignmentId), settlementKey(assignment.assignmentId), ), - this.redis.eval( - script, - 1, - lockKey(assignment.workerId), - 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])', + 'end', + 'return 0', + ].join('\n'); + await this.redis.eval(script, 1, lockKey(workerId), assignmentId); + } } diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index 671d1850..a18788a8 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -29,6 +29,7 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { 'No bridge worker is configured', ); } + const sessionResultFinalizer = ctx.sessionResultFinalizer; try { const settlement = await this.store.dispatch({ workerId: this.workerId, @@ -37,6 +38,15 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { 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( 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', From b68f28dc5dcb7d7773e5642126bf6257b3ba220e Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 07:49:49 -0400 Subject: [PATCH 04/29] fix: harden remote bridge assignment lifecycle --- docs/remote-bridge/README.md | 7 ++ packages/code/README.md | 10 +- packages/code/src/worker.test.ts | 198 ++++++++++++++++++++++++++++++- packages/code/src/worker.ts | 117 ++++++++++++++++-- service/src/bridge/store.test.ts | 74 ++++++++++++ service/src/bridge/store.ts | 25 ++-- 6 files changed, 402 insertions(+), 29 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index f5632d5d..a8d50e2a 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -67,6 +67,13 @@ execution. - 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; diff --git a/packages/code/README.md b/packages/code/README.md index a79a1709..edb8de2a 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -32,12 +32,20 @@ Optional environment variables: 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. + 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/src/worker.test.ts b/packages/code/src/worker.test.ts index 03061d91..58be72d5 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -1,6 +1,6 @@ import assert from 'node:assert/strict'; import test from 'node:test'; -import { BridgeWorker } from './worker.js'; +import { BridgeWorker, BridgeWorkspaceQuarantinedError } from './worker.js'; import type { BridgeAssignment } from './protocol.js'; @@ -148,7 +148,11 @@ test('worker refreshes its registration during a long assignment', async () => { }); } return new Response( - JSON.stringify({ protocolVersion: 1, accepted: true, body: init?.body }), + JSON.stringify({ + protocolVersion: 1, + accepted: true, + body: init?.body, + }), { status: 200, headers: { 'Content-Type': 'application/json' } }, ); }; @@ -179,3 +183,193 @@ test('worker refreshes its registration during a long assignment', async () => { 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 surfaces a definite settlement rejection without quarantining', 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: {} }, + }), + (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, + ); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 31083947..3b942009 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -32,6 +32,7 @@ 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 { @@ -43,6 +44,16 @@ function errorMessage(value: object): string | undefined { 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; @@ -111,8 +122,11 @@ export class BridgeWorker { } catch (error) { if (signal?.aborted) return; if ( - error instanceof BridgeProtocolError && - (error.status === 401 || error.status === 403 || error.status === 409) + error instanceof BridgeWorkspaceQuarantinedError || + (error instanceof BridgeProtocolError && + (error.status === 401 || + error.status === 403 || + error.status === 409)) ) { throw error; } @@ -156,10 +170,11 @@ export class BridgeWorker { ); let settlement: BridgeSettlement; try { + const sandboxSessionId = this.sandboxSessionIdFor(assignment); const headers = { ...assignment.request.headers, - ...(assignment.runtimeSessionId - ? { 'X-Runtime-Session-Id': assignment.runtimeSessionId } + ...(sandboxSessionId + ? { 'X-Runtime-Session-Id': sandboxSessionId } : {}), }; const response = await this.fetchImpl( @@ -204,20 +219,39 @@ export class BridgeWorker { } clearTimeout(deadlineTimer); - heartbeatController.abort(); - await heartbeat; cancellationController.abort(); await cancellationWatcher; - signal?.removeEventListener('abort', abortExecution); - await this.request( - this.assignmentUrl(assignment, 'settle'), - settlement, - signal, - ); + try { + 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) return this.sandboxEndpoint; + 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) @@ -271,6 +305,63 @@ export class BridgeWorker { ); } + 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 + ) { + 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, diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index c35a8fd8..0db397f8 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -144,6 +144,80 @@ describe('RedisBridgeStore', () => { ).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('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('releases the worker lock when generation allocation fails', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 04fb8e29..3f34b765 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -13,7 +13,6 @@ 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; -const MAX_ASSIGNMENT_TTL_SECONDS = 10 * 60; export type CodeBridgeAssignment = BridgeAssignment; export type CodeBridgeSettlement = BridgeSettlement< @@ -73,8 +72,8 @@ function workspaceQuarantineKey( return `${PREFIX}:worker:${workerId}:workspace:${sessionHash}:quarantined`; } -function queueKey(workerId: string): string { - return `${PREFIX}:worker:${workerId}:assignments`; +function queueKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:assignments`; } function generationKey(workerId: string): string { @@ -102,13 +101,7 @@ function tokenHash(token: string): string { } function assignmentTtlSeconds(deadlineAtMs: number): number { - return Math.max( - 1, - Math.min( - MAX_ASSIGNMENT_TTL_SECONDS, - Math.ceil((deadlineAtMs - Date.now()) / 1000) + 30, - ), - ); + return Math.max(1, Math.ceil((deadlineAtMs - Date.now()) / 1000) + 30); } async function delay(ms: number, signal?: AbortSignal): Promise { @@ -255,8 +248,12 @@ export class RedisBridgeStore { 'EX', ttlSeconds, ); - transaction.rpush(queueKey(args.workerId), assignmentId); - transaction.expire(queueKey(args.workerId), ttlSeconds); + const assignmentQueueKey = queueKey( + args.workerId, + assignment.incarnationId, + ); + transaction.rpush(assignmentQueueKey, assignmentId); + transaction.expire(assignmentQueueKey, ttlSeconds); await transaction.exec(); const settlement = await this.waitForSettlement( assignment, @@ -294,7 +291,9 @@ export class RedisBridgeStore { ): Promise { const deadline = Date.now() + waitMs; while (signal?.aborted !== true && Date.now() < deadline) { - const assignmentId = await this.redis.lpop(queueKey(workerId)); + 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())), From 14dca6b76c0ddab47007b85d294bb545ac780ebf Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 08:13:39 -0400 Subject: [PATCH 05/29] fix: close remote bridge commit races --- packages/code/src/worker.test.ts | 50 +++++++ packages/code/src/worker.ts | 15 ++ service/src/bridge/router.ts | 7 +- service/src/bridge/store.test.ts | 223 +++++++++++++++++++++++++++++ service/src/bridge/store.ts | 176 +++++++++++++++++++---- service/src/lifecycle.ts | 4 + service/src/secure-startup.test.ts | 14 ++ service/src/secure-startup.ts | 28 ++-- 8 files changed, 477 insertions(+), 40 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 58be72d5..38821376 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -373,3 +373,53 @@ test('worker quarantines stateful reuse after settlement stays ambiguous', async 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); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 3b942009..afd43c2d 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -169,6 +169,8 @@ export class BridgeWorker { cancellationController.signal, ); let settlement: BridgeSettlement; + let ambiguousSandboxError: unknown; + let sandboxRejectedExecution = false; try { const sandboxSessionId = this.sandboxSessionIdFor(assignment); const headers = { @@ -192,6 +194,7 @@ export class BridgeWorker { 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}`, @@ -207,6 +210,12 @@ export class BridgeWorker { result: payload, }; } catch (error) { + if ( + assignment.runtimeSessionId != null && + !sandboxRejectedExecution + ) { + ambiguousSandboxError = error; + } settlement = { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, @@ -222,6 +231,12 @@ export class BridgeWorker { 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(); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index c69e1e03..64635308 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -55,7 +55,12 @@ function isRecord(value: unknown): value is Record { } function sendStoreError(error: BridgeStoreError, res: Response): void { - const status = error.code === 'ASSIGNMENT_NOT_FOUND' ? 404 : 409; + const status = + error.code === 'ASSIGNMENT_NOT_FOUND' + ? 404 + : error.code === 'WORKER_BUSY' + ? 503 + : 409; res.status(status).json({ error: error.message, code: error.code }); } diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 0db397f8..4149f1be 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -8,8 +8,12 @@ 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(); }); @@ -186,6 +190,121 @@ describe('RedisBridgeStore', () => { }); }); + 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('keeps assignment state through deadlines longer than ten minutes', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -218,6 +337,110 @@ describe('RedisBridgeStore', () => { }); }); + 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, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 3f34b765..f0e34e50 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -132,6 +132,7 @@ export class RedisBridgeStore { 'local current = redis.call(\'GET\', KEYS[4])', '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', @@ -142,11 +143,12 @@ export class RedisBridgeStore { const result = Number( await this.redis.eval( script, - 4, + 5, workerKey(registration.workerId), incarnationFenceKey(registration.workerId, registration.incarnationId), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), + lockKey(registration.workerId), registration.incarnationId, JSON.stringify(registration), String(this.workerTtlSeconds), @@ -165,6 +167,12 @@ export class RedisBridgeStore { '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: { @@ -178,7 +186,7 @@ export class RedisBridgeStore { settlement: CodeBridgeSettlement, ) => Promise; }): Promise { - const registration = await this.registration(args.workerId); + let registration = await this.registration(args.workerId); if (registration == null) { throw new BridgeStoreError( 'WORKER_OFFLINE', @@ -224,6 +232,7 @@ export class RedisBridgeStore { } let assignment: StoredAssignment | undefined; + let resultCommitted = false; try { const generation = await this.redis.incr(generationKey(args.workerId)); assignment = { @@ -241,30 +250,50 @@ export class RedisBridgeStore { headers: args.headers, }, }; - const transaction = this.redis.multi(); - transaction.set( - assignmentKey(assignmentId), - JSON.stringify(assignment), - 'EX', - ttlSeconds, - ); - const assignmentQueueKey = queueKey( - args.workerId, - assignment.incarnationId, - ); - transaction.rpush(assignmentQueueKey, assignmentId); - transaction.expire(assignmentQueueKey, ttlSeconds); - await transaction.exec(); + 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) return settlement; + if (args.finalize == null) { + resultCommitted = true; + return settlement; + } try { - return await args.finalize(settlement); + const result = await args.finalize(settlement); + resultCommitted = true; + return result; } catch (error) { - await this.quarantine(args.workerId, registration.incarnationId); + await this.quarantine(args.workerId, assignment.incarnationId); if (args.runtimeSessionId !== undefined) { await this.redis.set( workspaceQuarantineKey(args.workerId, args.runtimeSessionId), @@ -274,11 +303,16 @@ export class RedisBridgeStore { throw error; } } finally { - await this.cancel(assignmentId); - if (assignment == null) { - await this.releaseLock(args.workerId, assignmentId); + 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.cleanup(assignment); + await this.cleanupDispatch(args.workerId, assignmentId, assignment); } } } @@ -355,12 +389,27 @@ export class RedisBridgeStore { ); } const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); - await this.redis.set( - settlementKey(assignmentId), - JSON.stringify(settlement), - 'EX', - ttlSeconds, + 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( @@ -424,6 +473,21 @@ export class RedisBridgeStore { 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', @@ -434,6 +498,62 @@ export class RedisBridgeStore { 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])', + 'return 1', + ].join('\n'); + const result = await this.redis.eval( + script, + 3, + workerIncarnationKey(assignment.workerId), + assignmentKey(assignment.assignmentId), + queueKey(assignment.workerId, assignment.incarnationId), + assignment.incarnationId, + JSON.stringify(assignment), + String(ttlSeconds), + assignment.assignmentId, + ); + 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( 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/secure-startup.test.ts b/service/src/secure-startup.test.ts index 113b4fc9..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, @@ -328,6 +329,19 @@ describe('sandbox backend policy', () => { 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', () => { configureValidLambda(); env.PTC_MODE = 'blocking'; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index dc9cb4ca..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); @@ -119,17 +135,7 @@ export function validateSandboxBackendPolicy(): void { ); } if (env.SANDBOX_BACKEND === 'remote-bridge') { - 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)', - ); - } + validateApiBridgePolicy(); return; } if (env.SANDBOX_BACKEND !== 'lambda-microvm') return; From b5f6ff05c4a8a8a8d8eb9bdbd62d88ca3d563783 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 08:19:02 -0400 Subject: [PATCH 06/29] fix: surface shutdown workspace quarantine --- packages/code/src/worker.test.ts | 72 ++++++++++++++++++++++++++++++++ packages/code/src/worker.ts | 10 ++--- 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 38821376..c2396569 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -423,3 +423,75 @@ test('worker quarantines a stateful workspace after the sandbox request aborts', ); 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 index afd43c2d..561dfc68 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -120,13 +120,13 @@ export class BridgeWorker { if (!assignment) continue; await this.executeAndSettle(assignment, signal); } catch (error) { + if (error instanceof BridgeWorkspaceQuarantinedError) { + throw error; + } if (signal?.aborted) return; if ( - error instanceof BridgeWorkspaceQuarantinedError || - (error instanceof BridgeProtocolError && - (error.status === 401 || - error.status === 403 || - error.status === 409)) + error instanceof BridgeProtocolError && + (error.status === 401 || error.status === 403 || error.status === 409) ) { throw error; } From 6c19e5c82952bd5eaa9f6c14d78228ae2d7aea4a Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 08:37:40 -0400 Subject: [PATCH 07/29] fix: fail closed across bridge lifecycle gaps --- packages/code/src/worker.test.ts | 44 ++++++- packages/code/src/worker.ts | 9 ++ service/src/bridge/router.ts | 193 +++++++++++++++++-------------- service/src/bridge/store.test.ts | 85 ++++++++++++++ service/src/bridge/store.ts | 110 +++++++++++++----- 5 files changed, 323 insertions(+), 118 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index c2396569..30f07993 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -239,7 +239,7 @@ test('worker routes a hintless assignment to an ephemeral template session', asy assert.equal(runtimeSessionHeader, 'assignment-hintless-assignment'); }); -test('worker surfaces a definite settlement rejection without quarantining', async () => { +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: [] }), { @@ -278,6 +278,48 @@ test('worker surfaces a definite settlement rejection without quarantining', asy 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' && diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 561dfc68..263c58c5 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -353,6 +353,15 @@ export class BridgeWorker { 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(); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 64635308..8653bade 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -1,7 +1,7 @@ import { timingSafeEqual } from 'crypto'; import { Router } from 'express'; -import type { NextFunction, Request, Response } from 'express'; +import type { NextFunction, Request, RequestHandler, Response } from 'express'; import type { BridgeWorkerRegistration } from '../../../packages/code/src/protocol'; import type { CodeBridgeSettlement } from './store'; @@ -54,6 +54,14 @@ 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' @@ -90,97 +98,108 @@ function isSettlement(value: unknown): value is CodeBridgeSettlement { const router = Router(); router.use(bridgeAuth); -router.post('/workers/register', async (req: Request, res: Response) => { - 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); +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; } - 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', async (req: Request, res: Response) => { - 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', + 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, }); - 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); + }), +); + +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; } - throw error; - } -}); + 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', - async (req, res) => { + asyncRoute(async (req, res) => { const settlement = req.body as unknown; if (!isSettlement(settlement)) { res.status(400).json({ error: 'Invalid bridge settlement' }); @@ -203,12 +222,12 @@ router.post( } throw error; } - }, + }), ); router.post( '/workers/:workerId/assignments/:assignmentId/cancellation', - async (req, res) => { + asyncRoute(async (req, res) => { const body = isRecord(req.body) ? req.body : {}; if ( body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || @@ -223,7 +242,7 @@ router.post( 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 index 4149f1be..6a4cf77b 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -1,4 +1,5 @@ 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'; @@ -305,6 +306,90 @@ describe('RedisBridgeStore', () => { ).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, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index f0e34e50..e0b09b8a 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -84,6 +84,10 @@ 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}`; } @@ -107,15 +111,15 @@ function assignmentTtlSeconds(deadlineAtMs: number): number { async function delay(ms: number, signal?: AbortSignal): Promise { if (signal?.aborted === true) return; await new Promise((resolve) => { - const timer = setTimeout(resolve, ms); - signal?.addEventListener( - 'abort', - () => { - clearTimeout(timer); - resolve(); - }, - { once: true }, - ); + const onAbort = (): void => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal?.addEventListener('abort', onAbort, { once: true }); }); } @@ -130,6 +134,10 @@ export class RedisBridgeStore { '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', @@ -143,12 +151,13 @@ export class RedisBridgeStore { const result = Number( await this.redis.eval( script, - 5, + 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), @@ -217,14 +226,13 @@ export class RedisBridgeStore { const assignmentId = randomBytes(18).toString('base64url'); const leaseToken = randomBytes(32).toString('base64url'); const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs); - const locked = await this.redis.set( - lockKey(args.workerId), + const locked = await this.acquireLock( + args.workerId, assignmentId, - 'PX', - ttlSeconds * 1000, - 'NX', + registration.incarnationId, + ttlSeconds, ); - if (locked !== 'OK') { + if (!locked) { throw new BridgeStoreError( 'WORKER_BUSY', `Bridge worker ${args.workerId} is busy`, @@ -293,13 +301,11 @@ export class RedisBridgeStore { resultCommitted = true; return result; } catch (error) { - await this.quarantine(args.workerId, assignment.incarnationId); - if (args.runtimeSessionId !== undefined) { - await this.redis.set( - workspaceQuarantineKey(args.workerId, args.runtimeSessionId), - '1', - ); - } + await this.quarantine( + args.workerId, + assignment.incarnationId, + args.runtimeSessionId, + ); throw error; } } finally { @@ -430,21 +436,32 @@ export class RedisBridgeStore { return (await this.redis.exists(cancellationKey(assignmentId))) === 1; } - async quarantine(workerId: string, incarnationId: string): Promise { + 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'); - await this.redis.eval( - script, - 3, + 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, ); } @@ -507,18 +524,45 @@ export class RedisBridgeStore { '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, - 3, + 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; } @@ -570,10 +614,16 @@ export class RedisBridgeStore { ): Promise { const script = [ 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', - ' return redis.call(\'DEL\', KEYS[1])', + ' return redis.call(\'DEL\', KEYS[1], KEYS[2])', 'end', 'return 0', ].join('\n'); - await this.redis.eval(script, 1, lockKey(workerId), assignmentId); + await this.redis.eval( + script, + 2, + lockKey(workerId), + lockIncarnationKey(workerId), + assignmentId, + ); } } From b970aba16e0e4ae7f01f4170c026446de09ddd85 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 08:57:43 -0400 Subject: [PATCH 08/29] fix: persist stateful settlement commit barriers --- docs/remote-bridge/README.md | 3 ++ packages/code/src/worker.test.ts | 76 ++++++++++++++++++++++++++++++ packages/code/src/worker.ts | 50 ++++++++++++++++---- service/src/bridge/store.test.ts | 67 +++++++++++++++++++++++++- service/src/bridge/store.ts | 64 +++++++++++++++++++++---- service/src/secure-startup.test.ts | 14 ++++++ service/src/secure-startup.ts | 1 + 7 files changed, 256 insertions(+), 19 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index a8d50e2a..3c6e7ba9 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -74,6 +74,9 @@ execution. - 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. +- A fulfilled stateful settlement creates a durable pending-workspace marker + before Code API acknowledges it. Result finalization clears that marker; a + worker-process crash leaves it in place so later reuse fails closed. - 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; diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 30f07993..c2f03c38 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -537,3 +537,79 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async ); assert.equal(executeStarted, true); }); + +test('worker bounds a stalled lease transport beyond its long poll', async () => { + 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'], + }, + leaseWaitMs: 10, + leaseTransportGraceMs: 20, + fetchImpl: async (_input, init) => + await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }), + }); + + await assert.rejects(worker.lease(), { name: 'AbortError' }); +}); + +test('worker quarantines an explicitly dirty stateful sandbox response', async () => { + let settlementAttempted = false; + 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: async (input) => { + if (String(input).endsWith('/execute')) { + return new Response( + JSON.stringify({ + error: 'session_workspace_dirty', + message: 'restore required', + }), + { status: 409, headers: { 'Content-Type': 'application/json' } }, + ); + } + settlementAttempted = true; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'dirty-execution', + 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, + ); + assert.equal(settlementAttempted, false); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 263c58c5..0ba49e8c 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -22,6 +22,7 @@ export interface BridgeWorkerOptions { sandboxEndpoint: string; capabilities: BridgeWorkerCapabilities; leaseWaitMs?: number; + leaseTransportGraceMs?: number; reconnectDelayMs?: number; fetchImpl?: typeof fetch; onError?: (error: unknown) => void; @@ -29,6 +30,8 @@ export interface BridgeWorkerOptions { } const DEFAULT_LEASE_WAIT_MS = 25_000; +const MAX_LEASE_WAIT_MS = 30_000; +const DEFAULT_LEASE_TRANSPORT_GRACE_MS = 5_000; const DEFAULT_RECONNECT_DELAY_MS = 1_000; const DEFAULT_REGISTRATION_TTL_MS = 60_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; @@ -92,15 +95,41 @@ export class BridgeWorker { } 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, + const waitMs = Math.min( + MAX_LEASE_WAIT_MS, + Math.max(0, this.options.leaseWaitMs ?? DEFAULT_LEASE_WAIT_MS), ); + const leaseController = new AbortController(); + const abortLease = (): void => leaseController.abort(); + if (signal?.aborted) { + abortLease(); + } else { + signal?.addEventListener('abort', abortLease, { once: true }); + } + const timeout = setTimeout( + abortLease, + waitMs + + Math.max( + 0, + this.options.leaseTransportGraceMs ?? + DEFAULT_LEASE_TRANSPORT_GRACE_MS, + ), + ); + let response: BridgeLeaseResponse; + try { + response = await this.request( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/lease`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + waitMs, + incarnationId: this.incarnationId, + }, + leaseController.signal, + ); + } finally { + clearTimeout(timeout); + signal?.removeEventListener('abort', abortLease); + } if ( response.assignment != null && response.assignment.incarnationId !== this.incarnationId @@ -194,7 +223,8 @@ export class BridgeWorker { const payload = (await response.json()) as object; if (heartbeatError != null) throw heartbeatError; if (!response.ok) { - sandboxRejectedExecution = true; + sandboxRejectedExecution = + errorMessage(payload) !== 'session_workspace_dirty'; throw new BridgeProtocolError( errorMessage(payload) ?? `Sandbox rejected execution with HTTP ${response.status}`, @@ -392,7 +422,7 @@ export class BridgeWorker { signal: AbortSignal, ): Promise { while (!signal.aborted && !executionController.signal.aborted) { - await new Promise((resolve) => setTimeout(resolve, 500)); + await this.delay(500, signal); if (signal.aborted || executionController.signal.aborted) return; try { const response = await this.request<{ cancelled: boolean }>( diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 6a4cf77b..4065b2f7 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -526,6 +526,71 @@ describe('RedisBridgeStore', () => { redis.del = originalDel as Redis['del']; }); + test('holds a durable workspace marker until finalization commits', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'commit-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + let releaseFinalizer!: () => void; + const finalizerGate = new Promise((resolve) => { + releaseFinalizer = resolve; + }); + let finalizerStarted!: () => void; + const started = new Promise((resolve) => { + finalizerStarted = resolve; + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'commit-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-commit', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + finalize: async (settlement) => { + finalizerStarted(); + await finalizerGate; + return settlement; + }, + }); + const assignment = await store.lease( + 'commit-worker', + incarnationId, + 1_000, + ); + await store.settle('commit-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-commit', + files: [], + }, + }); + await started; + const [pendingMarker] = await redis.keys( + 'codeapi:bridge:v1:worker:commit-worker:workspace:*:quarantined', + ); + expect(pendingMarker).toBeDefined(); + expect(await redis.get(pendingMarker)).toBe( + assignment?.assignmentId ?? null, + ); + + releaseFinalizer(); + await expect(completion).resolves.toMatchObject({ status: 'fulfilled' }); + expect(await redis.exists(pendingMarker)).toBe(0); + }); + test('releases the worker lock when generation allocation fails', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -602,7 +667,7 @@ describe('RedisBridgeStore', () => { deadlineAtMs: Date.now() + 1_000, signal: controller.signal, }), - ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); throw new Error('restore failed'); }, }); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index e0b09b8a..ef6c6bb3 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -292,12 +292,12 @@ export class RedisBridgeStore { args.deadlineAtMs, args.signal, ); - if (args.finalize == null) { - resultCommitted = true; - return settlement; - } try { - const result = await args.finalize(settlement); + const result = + args.finalize == null + ? settlement + : await args.finalize(settlement); + await this.commitPendingWorkspace(assignment, settlement); resultCommitted = true; return result; } catch (error) { @@ -395,19 +395,32 @@ export class RedisBridgeStore { ); } const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); + const settlementKeys = [ + assignmentKey(assignmentId), + settlementKey(assignmentId), + ]; + if ( + settlement.status === 'fulfilled' && + assignment.runtimeSessionId !== undefined + ) { + settlementKeys.push( + workspaceQuarantineKey(workerId, assignment.runtimeSessionId), + ); + } const script = [ 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return 0 end', 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', + 'if #KEYS == 3 then redis.call(\'SET\', KEYS[3], ARGV[3]) end', 'return 1', ].join('\n'); const accepted = Number( await this.redis.eval( script, - 2, - assignmentKey(assignmentId), - settlementKey(assignmentId), + settlementKeys.length, + ...settlementKeys, JSON.stringify(settlement), String(ttlSeconds), + assignmentId, ), ); if (accepted !== 1) { @@ -580,6 +593,41 @@ export class RedisBridgeStore { await this.cleanup(assignment); } + private async commitPendingWorkspace( + assignment: StoredAssignment, + settlement: CodeBridgeSettlement, + ): Promise { + if ( + assignment.runtimeSessionId === undefined || + settlement.status !== 'fulfilled' + ) { + return; + } + const script = [ + 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1])', + 'end', + 'return 0', + ].join('\n'); + const committed = Number( + await this.redis.eval( + script, + 1, + workspaceQuarantineKey( + assignment.workerId, + assignment.runtimeSessionId, + ), + assignment.assignmentId, + ), + ); + if (committed !== 1) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace commit marker was lost before finalization completed', + ); + } + } + private async cleanupWithRetry( workerId: string, assignmentId: string, diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 56e0fb23..c2e3c1f5 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -342,6 +342,20 @@ describe('sandbox backend policy', () => { expect(() => validateApiBridgePolicy()).not.toThrow(); }); + test('remote bridge requires a positive finite job timeout', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.BRIDGE_WORKER_ID = 'engineering-vm'; + env.BRIDGE_TOKEN = 'development-bridge-token'; + env.PTC_MODE = 'replay'; + + env.JOB_TIMEOUT = -1; + expect(() => validateApiBridgePolicy()).toThrow('JOB_TIMEOUT'); + env.JOB_TIMEOUT = Number.POSITIVE_INFINITY; + expect(() => validateApiBridgePolicy()).toThrow('JOB_TIMEOUT'); + env.JOB_TIMEOUT = 300_000; + expect(() => validateApiBridgePolicy()).not.toThrow(); + }); + test('rejects blocking PTC on the lambda backend', () => { configureValidLambda(); env.PTC_MODE = 'blocking'; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index c324cde4..437a6276 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -56,6 +56,7 @@ export function validateApiHardenedConfig(): void { /** Validate bridge credentials in every process that exposes bridge routes. */ export function validateApiBridgePolicy(): void { if (env.SANDBOX_BACKEND !== 'remote-bridge') return; + requireSafeWholeNumber('JOB_TIMEOUT', env.JOB_TIMEOUT, 1); requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); if (env.HARDENED_SANDBOX_MODE) { requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); From d4f150cc69e304ac21e98901a258284f719fcf17 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 09:16:43 -0400 Subject: [PATCH 09/29] fix: bound bridge control-plane timing --- packages/code/src/protocol.ts | 2 + packages/code/src/worker.test.ts | 69 +++++++++++++++++++++++++++++ packages/code/src/worker.ts | 76 +++++++++++++++++++++++++------- service/src/bridge/router.ts | 24 +++++++--- service/src/bridge/store.test.ts | 2 + service/src/bridge/store.ts | 5 ++- 6 files changed, 156 insertions(+), 22 deletions(-) diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index a79d4ff4..b19c262a 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -37,6 +37,8 @@ export interface BridgeAssignment { generation: number; leaseToken: string; expiresAt: string; + /** Server-calculated execution budget at lease time; avoids VM clock skew. */ + remainingMs?: number; runtimeSessionId?: string; request: BridgeSandboxRequest; } diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index c2f03c38..aea8e432 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -477,6 +477,7 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async generation: 1, leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(Date.now() + 5_000).toISOString(), + remainingMs: 5_000, runtimeSessionId: 'rt-user-1', request: { body: { language: 'bash' }, headers: {} }, }; @@ -613,3 +614,71 @@ test('worker quarantines an explicitly dirty stateful sandbox response', async ( ); assert.equal(settlementAttempted, false); }); + +test('worker bounds a stalled registration below its lease TTL', async () => { + 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'], + }, + registrationTransportTimeoutMs: 20, + fetchImpl: async (_input, init) => + await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }), + }); + + await assert.rejects(worker.register(), { name: 'AbortError' }); +}); + +test('worker uses the server-relative lease budget despite VM clock skew', async () => { + let settlementAttempted = false; + 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: async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + settlementAttempted = true; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'skewed-clock-assignment', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(0).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + assert.equal(settlementAttempted, true); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 0ba49e8c..c2cab051 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -23,6 +23,7 @@ export interface BridgeWorkerOptions { capabilities: BridgeWorkerCapabilities; leaseWaitMs?: number; leaseTransportGraceMs?: number; + registrationTransportTimeoutMs?: number; reconnectDelayMs?: number; fetchImpl?: typeof fetch; onError?: (error: unknown) => void; @@ -34,6 +35,7 @@ const MAX_LEASE_WAIT_MS = 30_000; const DEFAULT_LEASE_TRANSPORT_GRACE_MS = 5_000; const DEFAULT_RECONNECT_DELAY_MS = 1_000; const DEFAULT_REGISTRATION_TTL_MS = 60_000; +const DEFAULT_REGISTRATION_TRANSPORT_TIMEOUT_MS = 10_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; const SETTLEMENT_RETRY_DELAY_MS = 100; const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; @@ -75,16 +77,38 @@ export class BridgeWorker { 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, + const registrationController = new AbortController(); + const abortRegistration = (): void => registrationController.abort(); + if (signal?.aborted) { + abortRegistration(); + } else { + signal?.addEventListener('abort', abortRegistration, { once: true }); + } + const timeoutMs = Math.min( + Math.max(1, this.registrationTtlMs - 1), + Math.max( + 1, + this.options.registrationTransportTimeoutMs ?? + DEFAULT_REGISTRATION_TRANSPORT_TIMEOUT_MS, + ), ); + const timeout = setTimeout(abortRegistration, timeoutMs); + let registration: BridgeWorkerRegistrationResponse; + try { + registration = await this.request( + `${this.codeApiUrl}/bridge/workers/register`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: this.options.workerId, + incarnationId: this.incarnationId, + capabilities: this.options.capabilities, + }, + registrationController.signal, + ); + } finally { + clearTimeout(timeout); + signal?.removeEventListener('abort', abortRegistration); + } if (registration.incarnationId !== this.incarnationId) { throw new BridgeProtocolError( 'Code API registered a different worker incarnation', @@ -138,6 +162,15 @@ export class BridgeWorker { 'Code API leased an assignment for a different worker incarnation', ); } + if ( + response.assignment != null && + (!Number.isSafeInteger(response.assignment.remainingMs) || + (response.assignment.remainingMs ?? -1) < 0) + ) { + throw new BridgeProtocolError( + 'Code API leased an assignment without a valid server-relative deadline', + ); + } return response.assignment; } @@ -174,10 +207,8 @@ export class BridgeWorker { 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 deadlineDelay = this.assignmentRemainingMs(assignment); + const localDeadlineAtMs = Date.now() + deadlineDelay; const deadlineTimer = setTimeout( () => executionController.abort(), deadlineDelay, @@ -267,7 +298,12 @@ export class BridgeWorker { ambiguousSandboxError, ); } - await this.settleWithRetry(assignment, settlement, signal); + await this.settleWithRetry( + assignment, + settlement, + localDeadlineAtMs, + signal, + ); } finally { heartbeatController.abort(); await heartbeat; @@ -287,6 +323,16 @@ export class BridgeWorker { return undefined; } + private assignmentRemainingMs(assignment: BridgeAssignment): number { + if ( + Number.isSafeInteger(assignment.remainingMs) && + (assignment.remainingMs ?? -1) >= 0 + ) { + return assignment.remainingMs ?? 0; + } + return Math.max(0, Date.parse(assignment.expiresAt) - Date.now()); + } + private sandboxEndpointFor(assignment: BridgeAssignment): string { if (assignment.runtimeSessionId == null) { if (!this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER)) { @@ -353,9 +399,9 @@ export class BridgeWorker { private async settleWithRetry( assignment: BridgeAssignment, settlement: BridgeSettlement, + deadlineAtMs: number, signal?: AbortSignal, ): Promise { - const deadlineAtMs = Date.parse(assignment.expiresAt); const settlementController = new AbortController(); const abortSettlement = (): void => settlementController.abort(); signal?.addEventListener('abort', abortSettlement, { once: true }); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 8653bade..82156994 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -3,7 +3,7 @@ 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 type { CodeBridgeAssignment, CodeBridgeSettlement } from './store'; import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; import { connection } from '../queue'; @@ -181,11 +181,23 @@ router.post( return; } try { - const assignment = await bridgeStore.lease( - workerId, - body.incarnationId, - Math.min(requestedWait, MAX_LEASE_WAIT_MS), - ); + const leaseController = new AbortController(); + const abortLease = (): void => leaseController.abort(); + req.once('aborted', abortLease); + res.once('close', abortLease); + let assignment: CodeBridgeAssignment | undefined; + try { + assignment = await bridgeStore.lease( + workerId, + body.incarnationId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + leaseController.signal, + ); + } finally { + req.off('aborted', abortLease); + res.off('close', abortLease); + } + if (leaseController.signal.aborted) return; res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); } catch (error) { if (error instanceof BridgeStoreError) { diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 4065b2f7..bfce56ae 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -42,6 +42,8 @@ describe('RedisBridgeStore', () => { const assignment = await store.lease('vm-1', incarnationId, 1_000); expect(assignment).toBeDefined(); expect(assignment?.runtimeSessionId).toBe('rt-user-1'); + expect(assignment?.remainingMs).toBeGreaterThan(0); + expect(assignment?.remainingMs).toBeLessThanOrEqual(5_000); await store.settle('vm-1', assignment?.assignmentId ?? '', { protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index ef6c6bb3..2acdf8f9 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -353,7 +353,10 @@ export class RedisBridgeStore { } if (Date.parse(assignment.expiresAt) <= Date.now()) continue; const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; - return wireAssignment; + return { + ...wireAssignment, + remainingMs: Math.max(0, Date.parse(assignment.expiresAt) - Date.now()), + }; } return undefined; } From 38346c1e31255ba7938535eb412fe8f6691c13bd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 09:36:25 -0400 Subject: [PATCH 10/29] fix: recover abandoned bridge leases --- packages/code/src/worker.test.ts | 41 +++++++++++++ packages/code/src/worker.ts | 5 ++ service/src/bridge/router.ts | 22 ++++++- service/src/bridge/store.test.ts | 98 ++++++++++++++++++++++++++++++++ service/src/bridge/store.ts | 41 +++++++++++-- 5 files changed, 199 insertions(+), 8 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index aea8e432..1da96113 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -539,6 +539,47 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async assert.equal(executeStarted, true); }); +test('worker does not start execution after shutdown is already aborted', async () => { + const controller = new AbortController(); + controller.abort(new DOMException('shutdown', 'AbortError')); + let executeStarted = false; + 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: async () => { + executeStarted = true; + return new Response('{}', { status: 200 }); + }, + }); + + await assert.rejects( + worker.executeAndSettle( + { + protocolVersion: 1, + assignmentId: 'shutdown-before-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(), + remainingMs: 5_000, + request: { body: { language: 'bash' }, headers: {} }, + }, + controller.signal, + ), + { name: 'AbortError' }, + ); + assert.equal(executeStarted, false); +}); + test('worker bounds a stalled lease transport beyond its long poll', async () => { const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index c2cab051..1df457a5 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -204,6 +204,11 @@ export class BridgeWorker { assignment: BridgeAssignment, signal?: AbortSignal, ): Promise { + if (signal?.aborted === true) { + throw signal.reason instanceof Error + ? signal.reason + : new DOMException('aborted', 'AbortError'); + } const executionController = new AbortController(); const abortExecution = (): void => executionController.abort(); signal?.addEventListener('abort', abortExecution, { once: true }); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 82156994..06546ad8 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -193,12 +193,30 @@ router.post( Math.min(requestedWait, MAX_LEASE_WAIT_MS), leaseController.signal, ); + if (leaseController.signal.aborted) { + if (assignment != null) await bridgeStore.returnLease(assignment); + return; + } + const delivered = await new Promise((resolve) => { + const onFinish = (): void => { + res.off('close', onClose); + resolve(true); + }; + const onClose = (): void => { + res.off('finish', onFinish); + resolve(false); + }; + res.once('finish', onFinish); + res.once('close', onClose); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + }); + if (!delivered && assignment != null) { + await bridgeStore.returnLease(assignment); + } } finally { req.off('aborted', abortLease); res.off('close', abortLease); } - if (leaseController.signal.aborted) return; - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index bfce56ae..72db2816 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -11,10 +11,12 @@ const store = new RedisBridgeStore(redis); const incarnationId = 'incarnation-00000001'; const redisEval = redis.eval.bind(redis); const redisDel = redis.del.bind(redis); +const redisLpop = redis.lpop.bind(redis); afterEach(async () => { redis.eval = redisEval as Redis['eval']; redis.del = redisDel as Redis['del']; + redis.lpop = redisLpop as Redis['lpop']; await redis.flushall(); }); @@ -78,6 +80,46 @@ describe('RedisBridgeStore', () => { ).rejects.toMatchObject({ code: 'WORKER_OFFLINE' }); }); + test('returns a popped assignment when its lease request is aborted', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const dispatchController = new AbortController(); + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: dispatchController.signal, + }); + const leaseController = new AbortController(); + redis.lpop = (async (...args: Parameters) => { + const assignmentId = await redisLpop(...args); + if (assignmentId != null) leaseController.abort(); + return assignmentId; + }) as Redis['lpop']; + + await expect( + store.lease('vm-1', incarnationId, 1_000, leaseController.signal), + ).resolves.toBeUndefined(); + redis.lpop = redisLpop as Redis['lpop']; + + const recovered = await store.lease('vm-1', incarnationId, 1_000); + expect(recovered).toBeDefined(); + expect(recovered?.workerId).toBe('vm-1'); + dispatchController.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('rejects a stale lease token', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -723,4 +765,60 @@ describe('RedisBridgeStore', () => { }), ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); }); + + test('does not quarantine a stateless worker when finalization fails', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'stateless-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'stateless-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + finalize: async () => { + throw new Error('restore failed'); + }, + }); + const assignment = await store.lease( + 'stateless-worker', + incarnationId, + 1_000, + ); + await store.settle('stateless-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-1', + files: [], + }, + }); + + await expect(completion).rejects.toThrow('restore failed'); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'stateless-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).resolves.toBeUndefined(); + }); }); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 2acdf8f9..43b56679 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -123,6 +123,10 @@ async function delay(ms: number, signal?: AbortSignal): Promise { }); } +function signalAborted(signal?: AbortSignal): boolean { + return signal?.aborted === true; +} + export class RedisBridgeStore { constructor( private readonly redis: Redis, @@ -301,11 +305,13 @@ export class RedisBridgeStore { resultCommitted = true; return result; } catch (error) { - await this.quarantine( - args.workerId, - assignment.incarnationId, - args.runtimeSessionId, - ); + if (args.runtimeSessionId !== undefined) { + await this.quarantine( + args.workerId, + assignment.incarnationId, + args.runtimeSessionId, + ); + } throw error; } } finally { @@ -330,7 +336,7 @@ export class RedisBridgeStore { signal?: AbortSignal, ): Promise { const deadline = Date.now() + waitMs; - while (signal?.aborted !== true && Date.now() < deadline) { + while (!signalAborted(signal) && Date.now() < deadline) { const assignmentId = await this.redis.lpop( queueKey(workerId, incarnationId), ); @@ -344,6 +350,10 @@ export class RedisBridgeStore { const assignment = await this.readAssignment(assignmentId); if (assignment == null || assignment.workerId !== workerId) continue; if (assignment.incarnationId !== incarnationId) continue; + if (signalAborted(signal)) { + await this.returnLease(assignment); + return undefined; + } const registration = await this.registration(workerId); if (registration?.incarnationId !== incarnationId) { throw new BridgeStoreError( @@ -352,6 +362,10 @@ export class RedisBridgeStore { ); } if (Date.parse(assignment.expiresAt) <= Date.now()) continue; + if (signalAborted(signal)) { + await this.returnLease(assignment); + return undefined; + } const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; return { ...wireAssignment, @@ -361,6 +375,21 @@ export class RedisBridgeStore { return undefined; } + async returnLease(assignment: CodeBridgeAssignment): Promise { + await this.redis.eval( + [ + "if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end", + "redis.call('LREM', KEYS[2], 0, ARGV[1])", + "redis.call('LPUSH', KEYS[2], ARGV[1])", + 'return 1', + ].join('\n'), + 2, + assignmentKey(assignment.assignmentId), + queueKey(assignment.workerId, assignment.incarnationId), + assignment.assignmentId, + ); + } + async settle( workerId: string, assignmentId: string, From a8aca69f9c4cf007b7c001de4b1b9164ee653ee7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 09:58:35 -0400 Subject: [PATCH 11/29] fix: fence stateful bridge workspaces --- docs/remote-bridge/README.md | 8 +- packages/code/src/protocol.ts | 1 + packages/code/src/worker.test.ts | 83 +++++++++++++++++++++ packages/code/src/worker.ts | 11 ++- service/src/bridge/store.test.ts | 122 ++++++++++++++++++++++++++++++- service/src/bridge/store.ts | 71 ++++++++++++++---- 6 files changed, 275 insertions(+), 21 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 3c6e7ba9..1902e515 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -74,9 +74,11 @@ execution. - 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. -- A fulfilled stateful settlement creates a durable pending-workspace marker - before Code API acknowledges it. Result finalization clears that marker; a - worker-process crash leaves it in place so later reuse fails closed. +- Enqueueing stateful work atomically creates a durable in-flight workspace + marker. A definite rejection or successful result finalization clears it; + worker or VM loss leaves it in place so later reuse fails closed. Settlement + receipts outlive assignment cleanup briefly so retries are idempotent and + cannot recreate a cleared marker. - 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; diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index b19c262a..57d231a2 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -83,6 +83,7 @@ export class BridgeProtocolError extends Error { constructor( message: string, public readonly status?: number, + public readonly code?: string, ) { super(message); this.name = 'BridgeProtocolError'; diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 1da96113..805f8225 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -76,6 +76,89 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' }); }); +test('worker continues after an assignment-scoped settlement conflict', async () => { + const controller = new AbortController(); + let registrations = 0; + let leases = 0; + let observedError: unknown; + const assignment: BridgeAssignment = { + protocolVersion: 1, + assignmentId: 'expired-settlement', + 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(), + remainingMs: 5_000, + request: { body: { language: 'bash' }, headers: {} }, + }; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + registrations += 1; + if (registrations === 2) controller.abort(); + 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 (init?.signal?.aborted === true) { + throw new DOMException('aborted', 'AbortError'); + } + if (url.endsWith('/lease')) { + leases += 1; + return new Response( + JSON.stringify({ protocolVersion: 1, assignment }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + 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({ + error: 'Bridge assignment has expired', + code: 'ASSIGNMENT_EXPIRED', + }), + { 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'], + }, + reconnectDelayMs: 0, + fetchImpl, + onError: (error) => { + observedError = error; + }, + }); + + await worker.run(controller.signal); + assert.equal(registrations, 2); + assert.equal(leases, 1); + assert.equal( + observedError instanceof Error ? observedError.message : undefined, + 'Bridge assignment has expired', + ); +}); + test('worker aborts sandbox execution at the absolute assignment deadline', async () => { let settlement: Record | undefined; const fetchImpl: typeof fetch = async (input, init) => { diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 1df457a5..1b1116c8 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -49,6 +49,11 @@ function errorMessage(value: object): string | undefined { return undefined; } +function errorCode(value: object): string | undefined { + if ('code' in value && typeof value.code === 'string') return value.code; + return undefined; +} + export class BridgeWorkspaceQuarantinedError extends Error { constructor( message: string, @@ -188,7 +193,10 @@ export class BridgeWorker { if (signal?.aborted) return; if ( error instanceof BridgeProtocolError && - (error.status === 401 || error.status === 403 || error.status === 409) + (error.status === 401 || + error.status === 403 || + error.code === 'WORKER_FENCED' || + error.code === 'WORKER_QUARANTINED') ) { throw error; } @@ -518,6 +526,7 @@ export class BridgeWorker { errorMessage(payload) ?? `Bridge request failed with HTTP ${response.status}`, response.status, + errorCode(payload), ); } return payload as T; diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 72db2816..a3295fb2 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -608,19 +608,24 @@ describe('RedisBridgeStore', () => { incarnationId, 1_000, ); - await store.settle('commit-worker', assignment?.assignmentId ?? '', { + const settlement = { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment?.generation ?? 0, leaseToken: assignment?.leaseToken ?? '', incarnationId, - status: 'fulfilled', + status: 'fulfilled' as const, result: { language: 'bash', version: '5.2.0', session_id: 'run-commit', files: [], }, - }); + }; + await store.settle( + 'commit-worker', + assignment?.assignmentId ?? '', + settlement, + ); await started; const [pendingMarker] = await redis.keys( 'codeapi:bridge:v1:worker:commit-worker:workspace:*:quarantined', @@ -633,6 +638,117 @@ describe('RedisBridgeStore', () => { releaseFinalizer(); await expect(completion).resolves.toMatchObject({ status: 'fulfilled' }); expect(await redis.exists(pendingMarker)).toBe(0); + await expect( + store.settle( + 'commit-worker', + assignment?.assignmentId ?? '', + settlement, + ), + ).resolves.toBeUndefined(); + expect(await redis.exists(pendingMarker)).toBe(0); + }); + + test('keeps an in-flight workspace fenced when execution never settles', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'lost-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'lost-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-lost', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'lost-worker', + incarnationId, + 1_000, + ); + expect(assignment).toBeDefined(); + const [marker] = await redis.keys( + 'codeapi:bridge:v1:worker:lost-worker:workspace:*:quarantined', + ); + expect(await redis.get(marker)).toBe(assignment?.assignmentId ?? null); + + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'lost-worker', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + await expect( + store.dispatch({ + workerId: 'lost-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-lost', + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + }); + + test('clears an in-flight workspace marker after a definite rejection', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rejected-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'rejected-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-rejected', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'rejected-worker', + incarnationId, + 1_000, + ); + const [marker] = await redis.keys( + 'codeapi:bridge:v1:worker:rejected-worker:workspace:*:quarantined', + ); + expect(await redis.get(marker)).toBe(assignment?.assignmentId ?? null); + await store.settle( + 'rejected-worker', + assignment?.assignmentId ?? '', + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'rejected', + error: 'sandbox rejected before execution', + }, + ); + + await expect(completion).resolves.toMatchObject({ status: 'rejected' }); + expect(await redis.exists(marker)).toBe(0); }); test('releases the worker lock when generation allocation fails', async () => { diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 43b56679..afc4c415 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -395,6 +395,17 @@ export class RedisBridgeStore { assignmentId: string, settlement: CodeBridgeSettlement, ): Promise { + const serializedSettlement = JSON.stringify(settlement); + const existingSettlement = await this.redis.get( + settlementKey(assignmentId), + ); + if (existingSettlement === serializedSettlement) return; + if (existingSettlement != null) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment was already settled with a different result', + ); + } const assignment = await this.readAssignment(assignmentId); if (assignment == null) { throw new BridgeStoreError( @@ -431,18 +442,21 @@ export class RedisBridgeStore { assignmentKey(assignmentId), settlementKey(assignmentId), ]; - if ( - settlement.status === 'fulfilled' && - assignment.runtimeSessionId !== undefined - ) { + if (assignment.runtimeSessionId !== undefined) { settlementKeys.push( workspaceQuarantineKey(workerId, assignment.runtimeSessionId), ); } const script = [ + 'local existing = redis.call(\'GET\', KEYS[2])', + 'if existing then', + ' if existing == ARGV[1] then return 2 end', + ' return -1', + 'end', 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return 0 end', + 'if #KEYS == 3 and redis.call(\'GET\', KEYS[3]) ~= ARGV[3] then return -2 end', 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', - 'if #KEYS == 3 then redis.call(\'SET\', KEYS[3], ARGV[3]) end', + 'if #KEYS == 3 and ARGV[4] == \"rejected\" then redis.call(\'DEL\', KEYS[3]) end', 'return 1', ].join('\n'); const accepted = Number( @@ -450,12 +464,25 @@ export class RedisBridgeStore { script, settlementKeys.length, ...settlementKeys, - JSON.stringify(settlement), + serializedSettlement, String(ttlSeconds), assignmentId, + settlement.status, ), ); - if (accepted !== 1) { + if (accepted === -1) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment was already settled with a different result', + ); + } + if (accepted === -2) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace in-flight marker was lost before settlement', + ); + } + if (accepted !== 1 && accepted !== 2) { throw new BridgeStoreError( 'ASSIGNMENT_EXPIRED', 'Bridge assignment closed before settlement was committed', @@ -566,25 +593,44 @@ export class RedisBridgeStore { ): Promise { const script = [ 'if redis.call(\'GET\', KEYS[1]) ~= ARGV[1] then return 0 end', + 'if #KEYS == 5 and redis.call(\'EXISTS\', KEYS[5]) == 1 then return -1 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])', + 'if #KEYS == 5 then redis.call(\'SET\', KEYS[5], ARGV[4]) end', 'return 1', ].join('\n'); - const result = await this.redis.eval( - script, - 4, + const keys = [ workerIncarnationKey(assignment.workerId), assignmentKey(assignment.assignmentId), queueKey(assignment.workerId, assignment.incarnationId), lockIncarnationKey(assignment.workerId), + ]; + if (assignment.runtimeSessionId !== undefined) { + keys.push( + workspaceQuarantineKey( + assignment.workerId, + assignment.runtimeSessionId, + ), + ); + } + const result = await this.redis.eval( + script, + keys.length, + ...keys, assignment.incarnationId, JSON.stringify(assignment), String(ttlSeconds), assignment.assignmentId, String(ttlSeconds * 1000), ); + if (Number(result) === -1) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace already has incomplete stateful work', + ); + } return Number(result) === 1; } @@ -680,10 +726,7 @@ export class RedisBridgeStore { private async cleanup(assignment: StoredAssignment): Promise { await Promise.all([ - this.redis.del( - assignmentKey(assignment.assignmentId), - settlementKey(assignment.assignmentId), - ), + this.redis.del(assignmentKey(assignment.assignmentId)), this.releaseLock(assignment.workerId, assignment.assignmentId), ]); } From 0ec56ff281e170a6811256e12e54408767589159 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 10:16:52 -0400 Subject: [PATCH 12/29] fix: recover bridge lease read failures --- packages/code/src/worker.test.ts | 44 +++++++++++++++ packages/code/src/worker.ts | 11 ++++ service/src/bridge/store.test.ts | 46 ++++++++++++++++ service/src/bridge/store.ts | 92 +++++++++++++++++++++++--------- 4 files changed, 168 insertions(+), 25 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 805f8225..56c64373 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -663,6 +663,50 @@ test('worker does not start execution after shutdown is already aborted', async assert.equal(executeStarted, false); }); +test('worker does not start settlement after shutdown is already aborted', async () => { + const controller = new AbortController(); + let settlementAttempted = false; + 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: async (input) => { + if (String(input).endsWith('/execute')) { + controller.abort(new DOMException('shutdown', 'AbortError')); + throw new DOMException('aborted', 'AbortError'); + } + settlementAttempted = true; + return new Response('{}', { status: 200 }); + }, + }); + + await assert.rejects( + worker.executeAndSettle( + { + protocolVersion: 1, + assignmentId: 'shutdown-before-settlement', + 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(), + remainingMs: 5_000, + request: { body: { language: 'bash' }, headers: {} }, + }, + controller.signal, + ), + { name: 'AbortError' }, + ); + assert.equal(settlementAttempted, false); +}); + test('worker bounds a stalled lease transport beyond its long poll', async () => { const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 1b1116c8..d6299c2c 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -415,6 +415,17 @@ export class BridgeWorker { deadlineAtMs: number, signal?: AbortSignal, ): Promise { + if (signal?.aborted === true) { + if (assignment.runtimeSessionId != null) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined before settlement during shutdown`, + signal.reason, + ); + } + throw signal.reason instanceof Error + ? signal.reason + : new DOMException('aborted', 'AbortError'); + } const settlementController = new AbortController(); const abortSettlement = (): void => settlementController.abort(); signal?.addEventListener('abort', abortSettlement, { once: true }); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index a3295fb2..9f717818 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -12,11 +12,13 @@ const incarnationId = 'incarnation-00000001'; const redisEval = redis.eval.bind(redis); const redisDel = redis.del.bind(redis); const redisLpop = redis.lpop.bind(redis); +const redisGet = redis.get.bind(redis); afterEach(async () => { redis.eval = redisEval as Redis['eval']; redis.del = redisDel as Redis['del']; redis.lpop = redisLpop as Redis['lpop']; + redis.get = redisGet as Redis['get']; await redis.flushall(); }); @@ -120,6 +122,50 @@ describe('RedisBridgeStore', () => { }); }); + test('returns a popped assignment after a transient Redis read failure', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: false, + 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, + }); + let failAssignmentRead = true; + redis.get = (async (key: string) => { + if ( + failAssignmentRead && + key.includes(':assignment:') && + !key.endsWith(':settlement') + ) { + failAssignmentRead = false; + throw new Error('redis read failed'); + } + return await redisGet(key); + }) as Redis['get']; + + await expect(store.lease('vm-1', incarnationId, 1_000)).rejects.toThrow( + 'redis read failed', + ); + redis.get = redisGet as Redis['get']; + const recovered = await store.lease('vm-1', incarnationId, 1_000); + expect(recovered).toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('rejects a stale lease token', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index afc4c415..04ea8357 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -347,35 +347,59 @@ export class RedisBridgeStore { ); continue; } - const assignment = await this.readAssignment(assignmentId); - if (assignment == null || assignment.workerId !== workerId) continue; - if (assignment.incarnationId !== incarnationId) continue; - if (signalAborted(signal)) { - await this.returnLease(assignment); - return undefined; - } - const registration = await this.registration(workerId); - if (registration?.incarnationId !== incarnationId) { - throw new BridgeStoreError( - 'WORKER_FENCED', - 'Bridge worker incarnation was replaced', + try { + const assignment = await this.readAssignment(assignmentId); + if (assignment == null || assignment.workerId !== workerId) continue; + if (assignment.incarnationId !== incarnationId) continue; + if (signalAborted(signal)) { + await this.returnLease(assignment); + return undefined; + } + 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; + if (signalAborted(signal)) { + await this.returnLease(assignment); + return undefined; + } + const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; + return { + ...wireAssignment, + remainingMs: Math.max( + 0, + Date.parse(assignment.expiresAt) - Date.now(), + ), + }; + } catch (error) { + await this.returnLeaseByIdWithRetry( + workerId, + incarnationId, + assignmentId, ); + throw error; } - if (Date.parse(assignment.expiresAt) <= Date.now()) continue; - if (signalAborted(signal)) { - await this.returnLease(assignment); - return undefined; - } - const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; - return { - ...wireAssignment, - remainingMs: Math.max(0, Date.parse(assignment.expiresAt) - Date.now()), - }; } return undefined; } async returnLease(assignment: CodeBridgeAssignment): Promise { + await this.returnLeaseById( + assignment.workerId, + assignment.incarnationId, + assignment.assignmentId, + ); + } + + private async returnLeaseById( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { await this.redis.eval( [ "if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end", @@ -384,12 +408,30 @@ export class RedisBridgeStore { 'return 1', ].join('\n'), 2, - assignmentKey(assignment.assignmentId), - queueKey(assignment.workerId, assignment.incarnationId), - assignment.assignmentId, + assignmentKey(assignmentId), + queueKey(workerId, incarnationId), + assignmentId, ); } + private async returnLeaseByIdWithRetry( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { + let lastError: unknown; + for (let attempt = 0; attempt < 3; attempt += 1) { + try { + await this.returnLeaseById(workerId, incarnationId, assignmentId); + return; + } catch (error) { + lastError = error; + await delay(25); + } + } + throw lastError; + } + async settle( workerId: string, assignmentId: string, From 81a1061b4018a0f4b8c5c796d8fb25e9971c3606 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 10:34:47 -0400 Subject: [PATCH 13/29] fix: add safe workspace fence recovery --- docs/remote-bridge/README.md | 5 ++ packages/code/README.md | 5 ++ packages/code/src/cli.ts | 24 +++++++- packages/code/src/worker.test.ts | 74 ++++++++++++++++++++++++ packages/code/src/worker.ts | 24 +++++++- service/src/bridge/router.ts | 42 ++++++++++++++ service/src/bridge/store.test.ts | 96 ++++++++++++++++++++++++++++++++ service/src/bridge/store.ts | 49 ++++++++++++++++ 8 files changed, 317 insertions(+), 2 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 1902e515..125608f1 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -79,6 +79,11 @@ execution. worker or VM loss leaves it in place so later reuse fails closed. Settlement receipts outlive assignment cleanup briefly so retries are idempotent and cannot recreate a cleared marker. +- To recover a fenced session, first stop the worker and discard/reset that + session's local sandbox workspace. Restart registration, then run + `librechat-code reset-workspace ` with the same worker + configuration. Code API refuses the acknowledgement while work is active or + when it is not made by the currently registered incarnation. - 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; diff --git a/packages/code/README.md b/packages/code/README.md index edb8de2a..429350e5 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -46,6 +46,11 @@ 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. +After discarding or resetting that session's local runner, acknowledge recovery +with `librechat-code reset-workspace `. The command uses the +configured worker credentials, registers a fresh incarnation, and only clears +the server fence when no assignment is active. + 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/src/cli.ts b/packages/code/src/cli.ts index 411dda7f..a10e2ed4 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -49,7 +49,29 @@ const worker = new BridgeWorker({ }, }); -worker.run(controller.signal).catch((error: Error) => { +async function main(): Promise { + const command = process.argv[2]; + if (command === 'reset-workspace') { + const runtimeSessionId = process.argv[3]?.trim(); + if (!runtimeSessionId) { + throw new Error( + 'Usage: librechat-code reset-workspace ', + ); + } + await worker.register(controller.signal); + await worker.resetWorkspace(runtimeSessionId, controller.signal); + process.stdout.write( + `librechat-code: reset acknowledged for ${runtimeSessionId}\n`, + ); + return; + } + if (command != null) { + throw new Error(`Unknown command: ${command}`); + } + await worker.run(controller.signal); +} + +main().catch((error: Error) => { process.stderr.write(`librechat-code: ${error.message}\n`); process.exitCode = 1; }); diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 56c64373..5051cc4a 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -76,6 +76,38 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' }); }); +test('worker acknowledges a discarded workspace through the reset endpoint', async () => { + let requestBody: Record | undefined; + 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: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + assert.match(String(input), /workers\/vm-1\/workspaces\/reset$/); + requestBody = JSON.parse(String(init?.body)) as Record; + return new Response(JSON.stringify({ protocolVersion: 1, reset: true }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + }, + }); + + await worker.resetWorkspace('rt-user-1'); + assert.deepEqual(requestBody, { + protocolVersion: 1, + incarnationId: 'incarnation-00000001', + runtimeSessionId: 'rt-user-1', + confirmDiscarded: true, + }); +}); + test('worker continues after an assignment-scoped settlement conflict', async () => { const controller = new AbortController(); let registrations = 0; @@ -499,6 +531,48 @@ test('worker quarantines stateful reuse after settlement stays ambiguous', async ); }); +test('worker keeps a definite stateful rejection nonfatal when settlement is ambiguous', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ error: 'syntax_error' }), { + status: 400, + 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: 'rejected-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: {} }, + }), + (error: unknown) => + error instanceof TypeError && + !(error instanceof BridgeWorkspaceQuarantinedError), + ); +}); + test('worker quarantines a stateful workspace after the sandbox request aborts', async () => { let settlementAttempted = false; const fetchImpl: typeof fetch = async (input, init) => { diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index d6299c2c..cd933c56 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -123,6 +123,25 @@ export class BridgeWorker { return registration; } + async resetWorkspace( + runtimeSessionId: string, + signal?: AbortSignal, + ): Promise { + if (runtimeSessionId.trim().length === 0) { + throw new BridgeProtocolError('Runtime session ID is required'); + } + await this.request( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + runtimeSessionId, + confirmDiscarded: true, + }, + signal, + ); + } + async lease(signal?: AbortSignal): Promise { const waitMs = Math.min( MAX_LEASE_WAIT_MS, @@ -476,7 +495,10 @@ export class BridgeWorker { clearTimeout(deadlineTimer); signal?.removeEventListener('abort', abortSettlement); } - if (assignment.runtimeSessionId != null) { + if ( + assignment.runtimeSessionId != null && + settlement.status === 'fulfilled' + ) { throw new BridgeWorkspaceQuarantinedError( `Stateful workspace ${assignment.runtimeSessionId} was quarantined after ambiguous settlement delivery`, lastError, diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 06546ad8..ed9f6b60 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -158,6 +158,48 @@ router.post( }), ); +router.post( + '/workers/:workerId/workspaces/reset', + asyncRoute(async (req, res) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + typeof body.runtimeSessionId !== 'string' || + body.runtimeSessionId.trim().length === 0 || + body.runtimeSessionId.length > 512 || + body.confirmDiscarded !== true + ) { + res.status(400).json({ + error: 'Workspace reset requires confirmation of local discard', + }); + 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 { + await bridgeStore.resetWorkspace( + workerId, + body.incarnationId, + body.runtimeSessionId, + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), +); + router.post( '/workers/:workerId/lease', asyncRoute(async (req, res) => { diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 9f717818..52df42c1 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -82,6 +82,74 @@ describe('RedisBridgeStore', () => { ).rejects.toMatchObject({ code: 'WORKER_OFFLINE' }); }); + test('does not fence a workspace when dispatch is already aborted', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + controller.abort(); + await expect( + store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-aborted', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + expect( + await redis.keys( + 'codeapi:bridge:v1:worker:vm-1:workspace:*:quarantined', + ), + ).toHaveLength(0); + }); + + test('does not fence a workspace when dispatch aborts during lock acquisition', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + redis.eval = (async (...args: Parameters) => { + const result = await redisEval(...args); + if (String(args[0]).includes("EXISTS', KEYS[1]) == 1")) { + controller.abort(); + } + return result; + }) as Redis['eval']; + + await expect( + store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-aborted-lock', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + expect( + await redis.keys( + 'codeapi:bridge:v1:worker:vm-1:workspace:*:quarantined', + ), + ).toHaveLength(0); + expect(await redis.exists('codeapi:bridge:v1:worker:vm-1:lock')).toBe(0); + }); + test('returns a popped assignment when its lease request is aborted', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -724,6 +792,9 @@ describe('RedisBridgeStore', () => { 'codeapi:bridge:v1:worker:lost-worker:workspace:*:quarantined', ); expect(await redis.get(marker)).toBe(assignment?.assignmentId ?? null); + await expect( + store.resetWorkspace('lost-worker', incarnationId, 'rt-lost'), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); controller.abort(); await expect(completion).rejects.toMatchObject({ @@ -749,6 +820,31 @@ describe('RedisBridgeStore', () => { signal: new AbortController().signal, }), ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + + await store.resetWorkspace( + 'lost-worker', + 'incarnation-00000002', + 'rt-lost', + ); + const recoveredController = new AbortController(); + const recoveredCompletion = store.dispatch({ + workerId: 'lost-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-lost', + deadlineAtMs: Date.now() + 5_000, + signal: recoveredController.signal, + }); + const recoveredAssignment = await store.lease( + 'lost-worker', + 'incarnation-00000002', + 1_000, + ); + expect(recoveredAssignment).toBeDefined(); + recoveredController.abort(); + await expect(recoveredCompletion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); }); test('clears an in-flight workspace marker after a definite rejection', async () => { diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 04ea8357..e597fd8f 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -199,6 +199,7 @@ export class RedisBridgeStore { settlement: CodeBridgeSettlement, ) => Promise; }): Promise { + this.assertDispatchActive(args.signal, args.deadlineAtMs); let registration = await this.registration(args.workerId); if (registration == null) { throw new BridgeStoreError( @@ -246,6 +247,7 @@ export class RedisBridgeStore { let assignment: StoredAssignment | undefined; let resultCommitted = false; try { + this.assertDispatchActive(args.signal, args.deadlineAtMs); const generation = await this.redis.incr(generationKey(args.workerId)); assignment = { protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -264,6 +266,7 @@ export class RedisBridgeStore { }; let queued = false; for (let attempt = 0; attempt < 8 && !queued; attempt += 1) { + this.assertDispatchActive(args.signal, args.deadlineAtMs); assignment.incarnationId = registration.incarnationId; queued = await this.enqueueForActiveIncarnation(assignment, ttlSeconds); if (queued) break; @@ -580,6 +583,40 @@ export class RedisBridgeStore { ); } + async resetWorkspace( + workerId: string, + incarnationId: string, + runtimeSessionId: string, + ): Promise { + const result = Number( + await this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end", + "if redis.call('EXISTS', KEYS[2]) == 1 then return -2 end", + "redis.call('DEL', KEYS[3])", + 'return 1', + ].join('\n'), + 3, + workerIncarnationKey(workerId), + lockKey(workerId), + workspaceQuarantineKey(workerId, runtimeSessionId), + incarnationId, + ), + ); + if (result === -1) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Only the active bridge worker incarnation can reset a workspace', + ); + } + if (result === -2) { + throw new BridgeStoreError( + 'WORKER_BUSY', + 'Bridge workspace cannot be reset while worker execution is active', + ); + } + } + private async registration( workerId: string, ): Promise { @@ -587,6 +624,18 @@ export class RedisBridgeStore { return raw == null ? undefined : (JSON.parse(raw) as BridgeWorkerRegistration); } + private assertDispatchActive( + signal: AbortSignal, + deadlineAtMs: number, + ): void { + if (signal.aborted || Date.now() >= deadlineAtMs) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment ended before it could be delivered', + ); + } + } + private async readAssignment( assignmentId: string, ): Promise { From ea37f37e9674cff535d80378a019772668866bd6 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 10:53:30 -0400 Subject: [PATCH 14/29] fix: bound bridge liveness timers --- packages/code/src/worker.test.ts | 72 +++++++++++++++++++++++++++++++- packages/code/src/worker.ts | 42 +++++++++++++++++-- service/src/bridge/store.test.ts | 48 +++++++++++++++++++++ service/src/bridge/store.ts | 25 ++++++++++- 4 files changed, 180 insertions(+), 7 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 5051cc4a..6c80fa45 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -250,13 +250,13 @@ test('worker refreshes its registration during a long assignment', async () => { workerId: 'vm-1', incarnationId: 'incarnation-00000001', registeredAt: new Date().toISOString(), - leaseTtlMs: 50, + leaseTtlMs: 100, }), { status: 200, headers: { 'Content-Type': 'application/json' } }, ); } if (url.endsWith('/execute')) { - await new Promise((resolve) => setTimeout(resolve, 90)); + await new Promise((resolve) => setTimeout(resolve, 20)); return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { status: 200, headers: { 'Content-Type': 'application/json' }, @@ -285,6 +285,7 @@ test('worker refreshes its registration during a long assignment', async () => { fetchImpl, }); await worker.register(); + await new Promise((resolve) => setTimeout(resolve, 45)); await worker.executeAndSettle({ protocolVersion: 1, assignmentId: 'assignment-heartbeat', @@ -299,6 +300,73 @@ test('worker refreshes its registration during a long assignment', async () => { assert.ok(registrations >= 2); }); +test('worker continues cancellation polling after a stalled response', async () => { + let cancellationAttempts = 0; + let settlementAttempted = false; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + if (url.endsWith('/cancellation')) { + cancellationAttempts += 1; + if (cancellationAttempts === 1) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + return new Response(JSON.stringify({ cancelled: true }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + 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/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 5, + cancellationTransportTimeoutMs: 10, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'cancel-after-stall', + 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(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + assert.equal(cancellationAttempts, 2); + assert.equal(settlementAttempted, true); +}); + test('worker routes a hintless assignment to an ephemeral template session', async () => { let executeUrl = ''; let runtimeSessionHeader = ''; diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index cd933c56..2b3ed216 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -24,6 +24,8 @@ export interface BridgeWorkerOptions { leaseWaitMs?: number; leaseTransportGraceMs?: number; registrationTransportTimeoutMs?: number; + cancellationPollIntervalMs?: number; + cancellationTransportTimeoutMs?: number; reconnectDelayMs?: number; fetchImpl?: typeof fetch; onError?: (error: unknown) => void; @@ -36,6 +38,8 @@ const DEFAULT_LEASE_TRANSPORT_GRACE_MS = 5_000; const DEFAULT_RECONNECT_DELAY_MS = 1_000; const DEFAULT_REGISTRATION_TTL_MS = 60_000; const DEFAULT_REGISTRATION_TRANSPORT_TIMEOUT_MS = 10_000; +const DEFAULT_CANCELLATION_POLL_INTERVAL_MS = 500; +const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; const SETTLEMENT_RETRY_DELAY_MS = 100; const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; @@ -70,6 +74,7 @@ export class BridgeWorker { private readonly sandboxEndpoint: string; private readonly incarnationId: string; private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; + private lastRegisteredAtMs = 0; constructor(private readonly options: BridgeWorkerOptions) { this.fetchImpl = options.fetchImpl ?? fetch; @@ -120,6 +125,7 @@ export class BridgeWorker { ); } this.registrationTtlMs = registration.leaseTtlMs; + this.lastRegisteredAtMs = Date.now(); return registration; } @@ -245,6 +251,9 @@ export class BridgeWorker { () => executionController.abort(), deadlineDelay, ); + if (this.lastRegisteredAtMs === 0) { + this.lastRegisteredAtMs = Date.now(); + } const heartbeatController = new AbortController(); let heartbeatError: unknown; const heartbeat = this.maintainRegistration( @@ -394,10 +403,14 @@ export class BridgeWorker { executionController: AbortController, ): Promise { while (!signal.aborted && !executionController.signal.aborted) { + const heartbeatIntervalMs = Math.max( + MIN_REGISTRATION_HEARTBEAT_MS, + Math.floor(this.registrationTtlMs / 2), + ); await this.delay( Math.max( - MIN_REGISTRATION_HEARTBEAT_MS, - Math.floor(this.registrationTtlMs / 2), + 0, + this.lastRegisteredAtMs + heartbeatIntervalMs - Date.now(), ), signal, ); @@ -514,8 +527,26 @@ export class BridgeWorker { signal: AbortSignal, ): Promise { while (!signal.aborted && !executionController.signal.aborted) { - await this.delay(500, signal); + await this.delay( + Math.max( + 1, + this.options.cancellationPollIntervalMs ?? + DEFAULT_CANCELLATION_POLL_INTERVAL_MS, + ), + signal, + ); if (signal.aborted || executionController.signal.aborted) return; + const pollController = new AbortController(); + const abortPoll = (): void => pollController.abort(); + signal.addEventListener('abort', abortPoll, { once: true }); + const timeout = setTimeout( + abortPoll, + Math.max( + 1, + this.options.cancellationTransportTimeoutMs ?? + DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS, + ), + ); try { const response = await this.request<{ cancelled: boolean }>( this.assignmentUrl(assignment, 'cancellation'), @@ -523,7 +554,7 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, incarnationId: this.incarnationId, }, - signal, + pollController.signal, ); if (response.cancelled) { executionController.abort(); @@ -535,6 +566,9 @@ export class BridgeWorker { executionController.abort(); return; } + } finally { + clearTimeout(timeout); + signal.removeEventListener('abort', abortPoll); } } } diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 52df42c1..f6d7f52c 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -150,6 +150,54 @@ describe('RedisBridgeStore', () => { expect(await redis.exists('codeapi:bridge:v1:worker:vm-1:lock')).toBe(0); }); + test('clears a workspace fence when a queued assignment expires undelivered', 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-expired-queue', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + const queue = + 'codeapi:bridge:v1:worker:vm-1:incarnation:' + + `${incarnationId}:assignments`; + const assignmentId = await redis.lindex(queue, 0); + const assignmentKey = `codeapi:bridge:v1:assignment:${assignmentId}`; + const rawAssignment = await redis.get(assignmentKey); + const assignment = JSON.parse(rawAssignment ?? '{}') as Record< + string, + unknown + >; + assignment.expiresAt = new Date(0).toISOString(); + await redis.set(assignmentKey, JSON.stringify(assignment), 'EX', 30); + + await expect( + store.lease('vm-1', incarnationId, 100), + ).resolves.toBeUndefined(); + expect( + await redis.keys( + 'codeapi:bridge:v1:worker:vm-1:workspace:*:quarantined', + ), + ).toHaveLength(0); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('returns a popped assignment when its lease request is aborted', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index e597fd8f..ff4cdb86 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -365,7 +365,10 @@ export class RedisBridgeStore { 'Bridge worker incarnation was replaced', ); } - if (Date.parse(assignment.expiresAt) <= Date.now()) continue; + if (Date.parse(assignment.expiresAt) <= Date.now()) { + await this.clearUndeliveredWorkspaceFence(assignment); + continue; + } if (signalAborted(signal)) { await this.returnLease(assignment); return undefined; @@ -435,6 +438,26 @@ export class RedisBridgeStore { throw lastError; } + private async clearUndeliveredWorkspaceFence( + assignment: StoredAssignment, + ): Promise { + if (assignment.runtimeSessionId === undefined) return; + await this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) == ARGV[1] then", + " return redis.call('DEL', KEYS[1])", + 'end', + 'return 0', + ].join('\n'), + 1, + workspaceQuarantineKey( + assignment.workerId, + assignment.runtimeSessionId, + ), + assignment.assignmentId, + ); + } + async settle( workerId: string, assignmentId: string, From d23024f69e49cdcc3e3e31605cb3b9f4b044f2cb Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 11:14:35 -0400 Subject: [PATCH 15/29] fix: bound bridge cleanup recovery --- docs/remote-bridge/README.md | 10 +-- packages/code/README.md | 3 +- service/src/bridge/store.test.ts | 85 +++++++++++++++++++++-- service/src/bridge/store.ts | 112 ++++++++++++++++++++++++++++--- 4 files changed, 188 insertions(+), 22 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 125608f1..b67fc29b 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -79,11 +79,13 @@ execution. worker or VM loss leaves it in place so later reuse fails closed. Settlement receipts outlive assignment cleanup briefly so retries are idempotent and cannot recreate a cleared marker. -- To recover a fenced session, first stop the worker and discard/reset that - session's local sandbox workspace. Restart registration, then run +- To recover a fenced session, stop the normal worker process and discard/reset + that session's local sandbox workspace. While it remains stopped, run `librechat-code reset-workspace ` with the same worker - configuration. Code API refuses the acknowledgement while work is active or - when it is not made by the currently registered incarnation. + configuration; the command temporarily registers its own incarnation and + exits. Start the normal worker only after the reset command succeeds. Code API + refuses the acknowledgement while work is active or when it is not made by + the currently registered incarnation. - 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; diff --git a/packages/code/README.md b/packages/code/README.md index 429350e5..2439ab25 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -49,7 +49,8 @@ API did not commit. After discarding or resetting that session's local runner, acknowledge recovery with `librechat-code reset-workspace `. The command uses the configured worker credentials, registers a fresh incarnation, and only clears -the server fence when no assignment is active. +the server fence when no assignment is active. Run it while the normal worker +process is stopped, then restart the normal worker after the command exits. 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 diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index f6d7f52c..81454981 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -198,6 +198,43 @@ describe('RedisBridgeStore', () => { }); }); + test('clears a workspace fence when dispatch cancels before lease', 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-cancelled-queue', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + expect( + await redis.keys( + 'codeapi:bridge:v1:worker:vm-1:workspace:*:quarantined', + ), + ).toHaveLength(0); + expect( + await redis.llen( + `codeapi:bridge:v1:worker:vm-1:incarnation:${incarnationId}:assignments`, + ), + ).toBe(0); + }); + test('returns a popped assignment when its lease request is aborted', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -596,6 +633,37 @@ describe('RedisBridgeStore', () => { expect(getEventListeners(controller.signal, 'abort')).toHaveLength(0); }); + test('bounds a stalled Redis settlement poll command', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 20); + await timedStore.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'stalled-redis-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + redis.get = ((key: string) => { + if (key.endsWith(':settlement')) { + return new Promise(() => {}); + } + return redisGet(key); + }) as Redis['get']; + const controller = new AbortController(); + + await expect( + timedStore.dispatch({ + workerId: 'stalled-redis-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }), + ).rejects.toThrow('Bridge settlement poll timed out'); + }); + test('keeps assignment state through deadlines longer than ten minutes', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -703,13 +771,16 @@ describe('RedisBridgeStore', () => { 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']; + redis.eval = (async (...args: Parameters) => { + if (String(args[0]).includes("local queued = redis.call('LREM'")) { + cleanupAttempts += 1; + if (cleanupAttempts === 1) { + throw new Error('transient cleanup failure'); + } + } + return await redisEval(...args); + }) as Redis['eval']; await store.settle('cleanup-worker', assignment?.assignmentId ?? '', { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment?.generation ?? 0, @@ -729,7 +800,7 @@ describe('RedisBridgeStore', () => { result: { session_id: 'run-cleanup' }, }); expect(cleanupAttempts).toBeGreaterThanOrEqual(2); - redis.del = originalDel as Redis['del']; + redis.eval = redisEval as Redis['eval']; }); test('holds a durable workspace marker until finalization commits', async () => { diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index ff4cdb86..61aa3f7a 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -13,6 +13,7 @@ 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; +const DEFAULT_REDIS_COMMAND_TIMEOUT_MS = 1_000; export type CodeBridgeAssignment = BridgeAssignment; export type CodeBridgeSettlement = BridgeSettlement< @@ -127,10 +128,49 @@ function signalAborted(signal?: AbortSignal): boolean { return signal?.aborted === true; } +async function boundedCommand( + command: Promise, + timeoutMs: number, + label: string, + signal?: AbortSignal, +): Promise { + void command.catch(() => undefined); + return await new Promise((resolve, reject) => { + let settled = false; + const finish = (callback: () => void): void => { + if (settled) return; + settled = true; + clearTimeout(timer); + signal?.removeEventListener('abort', onAbort); + callback(); + }; + const onAbort = (): void => + finish(() => + reject( + signal?.reason instanceof Error + ? signal.reason + : new Error(`${label} aborted`), + ), + ); + const timer = setTimeout( + () => finish(() => reject(new Error(`${label} timed out`))), + timeoutMs, + ); + timer.unref?.(); + signal?.addEventListener('abort', onAbort, { once: true }); + if (signal?.aborted) onAbort(); + command.then( + (value) => finish(() => resolve(value)), + (error) => finish(() => reject(error)), + ); + }); +} + export class RedisBridgeStore { constructor( private readonly redis: Redis, private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, + private readonly redisCommandTimeoutMs = DEFAULT_REDIS_COMMAND_TIMEOUT_MS, ) {} async register(registration: BridgeWorkerRegistration): Promise { @@ -672,7 +712,15 @@ export class RedisBridgeStore { signal: AbortSignal, ): Promise { while (!signal.aborted && Date.now() < deadlineAtMs) { - const raw = await this.redis.get(settlementKey(assignment.assignmentId)); + const raw = await boundedCommand( + this.redis.get(settlementKey(assignment.assignmentId)), + Math.max( + 1, + Math.min(this.redisCommandTimeoutMs, deadlineAtMs - Date.now()), + ), + 'Bridge settlement poll', + signal, + ); if (raw != null) return JSON.parse(raw) as CodeBridgeSettlement; await delay(POLL_INTERVAL_MS, signal); } @@ -682,11 +730,15 @@ export class RedisBridgeStore { 'redis.call(\'DEL\', KEYS[1])', 'return nil', ].join('\n'); - const finalSettlement = await this.redis.eval( - closeScript, - 2, - assignmentKey(assignment.assignmentId), - settlementKey(assignment.assignmentId), + const finalSettlement = await boundedCommand( + this.redis.eval( + closeScript, + 2, + assignmentKey(assignment.assignmentId), + settlementKey(assignment.assignmentId), + ), + this.redisCommandTimeoutMs, + 'Bridge settlement close', ); if (finalSettlement != null) { return JSON.parse(String(finalSettlement)) as CodeBridgeSettlement; @@ -698,7 +750,11 @@ export class RedisBridgeStore { } private async cancel(assignmentId: string): Promise { - await this.redis.set(cancellationKey(assignmentId), '1', 'EX', 30); + await boundedCommand( + this.redis.set(cancellationKey(assignmentId), '1', 'EX', 30), + this.redisCommandTimeoutMs, + 'Bridge assignment cancellation', + ); } private async enqueueForActiveIncarnation( @@ -779,7 +835,11 @@ export class RedisBridgeStore { ): Promise { await this.cancel(assignmentId); if (assignment == null) { - await this.releaseLock(workerId, assignmentId); + await boundedCommand( + this.releaseLock(workerId, assignmentId), + this.redisCommandTimeoutMs, + 'Bridge assignment lock release', + ); return; } await this.cleanup(assignment); @@ -839,9 +899,41 @@ export class RedisBridgeStore { } private async cleanup(assignment: StoredAssignment): Promise { + const keys = [ + assignmentKey(assignment.assignmentId), + queueKey(assignment.workerId, assignment.incarnationId), + ]; + if (assignment.runtimeSessionId !== undefined) { + keys.push( + workspaceQuarantineKey( + assignment.workerId, + assignment.runtimeSessionId, + ), + ); + } + const cleanupScript = [ + "local queued = redis.call('LREM', KEYS[2], 0, ARGV[1])", + 'if queued > 0 and #KEYS == 3 and redis.call(\'GET\', KEYS[3]) == ARGV[1] then', + " redis.call('DEL', KEYS[3])", + 'end', + "return redis.call('DEL', KEYS[1])", + ].join('\n'); await Promise.all([ - this.redis.del(assignmentKey(assignment.assignmentId)), - this.releaseLock(assignment.workerId, assignment.assignmentId), + boundedCommand( + this.redis.eval( + cleanupScript, + keys.length, + ...keys, + assignment.assignmentId, + ), + this.redisCommandTimeoutMs, + 'Bridge assignment cleanup', + ), + boundedCommand( + this.releaseLock(assignment.workerId, assignment.assignmentId), + this.redisCommandTimeoutMs, + 'Bridge assignment lock release', + ), ]); } From 4ad81397c62b0f1baa366fe6b876a1638be08f2b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 11:43:34 -0400 Subject: [PATCH 16/29] fix: preserve bridge deadline fencing --- packages/code/src/protocol.ts | 2 + packages/code/src/worker.test.ts | 106 +++++++++++++++++++++++- packages/code/src/worker.ts | 38 ++++++++- service/src/bridge/router.ts | 7 +- service/src/bridge/store.test.ts | 72 +++++++++++++++++ service/src/bridge/store.ts | 133 +++++++++++++++++++++++-------- service/src/local-api.ts | 2 + 7 files changed, 321 insertions(+), 39 deletions(-) diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 57d231a2..eb9ec7d0 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -45,6 +45,8 @@ export interface BridgeAssignment { export interface BridgeLeaseResponse { protocolVersion: BridgeProtocolVersion; + /** Time spent handling the lease request on Code API, excluding transit. */ + serverElapsedMs?: number; assignment?: BridgeAssignment; } diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 6c80fa45..925be950 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -146,7 +146,7 @@ test('worker continues after an assignment-scoped settlement conflict', async () if (url.endsWith('/lease')) { leases += 1; return new Response( - JSON.stringify({ protocolVersion: 1, assignment }), + JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0, assignment }), { status: 200, headers: { 'Content-Type': 'application/json' } }, ); } @@ -621,6 +621,7 @@ test('worker keeps a definite stateful rejection nonfatal when settlement is amb runtimes: ['bash'], }, fetchImpl, + rejectionAckGraceMs: 0, }); await assert.rejects( @@ -641,6 +642,61 @@ test('worker keeps a definite stateful rejection nonfatal when settlement is amb ); }); +test('worker retries a known-clean rejection after shutdown until acknowledged', async () => { + const controller = new AbortController(); + let settlementAttempts = 0; + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ error: 'syntax_error' }), { + status: 400, + headers: { 'Content-Type': 'application/json' }, + }); + } + settlementAttempts += 1; + if (settlementAttempts === 1) { + controller.abort(); + 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'], + }, + rejectionAckGraceMs: 500, + fetchImpl, + }); + + await worker.executeAndSettle( + { + protocolVersion: 1, + assignmentId: 'late-clean-rejection', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 20).toISOString(), + remainingMs: 20, + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }, + controller.signal, + ); + assert.equal(controller.signal.aborted, true); + assert.equal(settlementAttempts, 2); +}); + test('worker quarantines a stateful workspace after the sandbox request aborts', async () => { let settlementAttempted = false; const fetchImpl: typeof fetch = async (input, init) => { @@ -722,7 +778,7 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async } if (url.endsWith('/lease')) { return new Response( - JSON.stringify({ protocolVersion: 1, assignment }), + JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0, assignment }), { status: 200, headers: { 'Content-Type': 'application/json' } }, ); } @@ -876,6 +932,52 @@ test('worker bounds a stalled lease transport beyond its long poll', async () => await assert.rejects(worker.lease(), { name: 'AbortError' }); }); +test('worker subtracts lease response transit from the server budget', async () => { + const originalNow = Date.now; + let now = 10_000; + Date.now = () => now; + try { + 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: async () => { + now += 50; + return new Response( + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 20, + assignment: { + protocolVersion: 1, + assignmentId: 'transit-budget', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(0).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + const assignment = await worker.lease(); + assert.equal(assignment?.remainingMs, 970); + } finally { + Date.now = originalNow; + } +}); + test('worker quarantines an explicitly dirty stateful sandbox response', async () => { let settlementAttempted = false; const worker = new BridgeWorker({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 2b3ed216..a47fadb6 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -26,6 +26,7 @@ export interface BridgeWorkerOptions { registrationTransportTimeoutMs?: number; cancellationPollIntervalMs?: number; cancellationTransportTimeoutMs?: number; + rejectionAckGraceMs?: number; reconnectDelayMs?: number; fetchImpl?: typeof fetch; onError?: (error: unknown) => void; @@ -42,6 +43,7 @@ const DEFAULT_CANCELLATION_POLL_INTERVAL_MS = 500; const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; const SETTLEMENT_RETRY_DELAY_MS = 100; +const REJECTION_ACK_GRACE_MS = 30_000; const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; function normalizedBaseUrl(value: string): string { @@ -170,6 +172,7 @@ export class BridgeWorker { ), ); let response: BridgeLeaseResponse; + const requestStartedAtMs = Date.now(); try { response = await this.request( `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/lease`, @@ -201,7 +204,26 @@ export class BridgeWorker { 'Code API leased an assignment without a valid server-relative deadline', ); } - return response.assignment; + if (response.assignment == null) return undefined; + if ( + !Number.isSafeInteger(response.serverElapsedMs) || + (response.serverElapsedMs ?? -1) < 0 + ) { + throw new BridgeProtocolError( + 'Code API leased an assignment without valid server timing', + ); + } + const transportElapsedMs = Math.max( + 0, + Date.now() - requestStartedAtMs - (response.serverElapsedMs ?? 0), + ); + return { + ...response.assignment, + remainingMs: Math.max( + 0, + (response.assignment.remainingMs ?? 0) - transportElapsedMs, + ), + }; } async run(signal?: AbortSignal): Promise { @@ -339,11 +361,21 @@ export class BridgeWorker { ambiguousSandboxError, ); } + const knownCleanStatefulRejection = + assignment.runtimeSessionId != null && + settlement.status === 'rejected' && + sandboxRejectedExecution; await this.settleWithRetry( assignment, settlement, - localDeadlineAtMs, - signal, + localDeadlineAtMs + + (knownCleanStatefulRejection + ? Math.max( + 0, + this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, + ) + : 0), + knownCleanStatefulRejection ? undefined : signal, ); } finally { heartbeatController.abort(); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index ed9f6b60..88142443 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -203,6 +203,7 @@ router.post( router.post( '/workers/:workerId/lease', asyncRoute(async (req, res) => { + const requestStartedAtMs = Date.now(); const workerId = req.params.workerId; const body = isRecord(req.body) ? req.body : {}; const requestedWait = Number(body.waitMs ?? 25_000); @@ -250,7 +251,11 @@ router.post( }; res.once('finish', onFinish); res.once('close', onClose); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), + assignment, + }); }); if (!delivered && assignment != null) { await bridgeStore.returnLease(assignment); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 81454981..f637ace4 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -664,6 +664,31 @@ describe('RedisBridgeStore', () => { ).rejects.toThrow('Bridge settlement poll timed out'); }); + test('bounds a stalled Redis dispatch preparation command', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 20); + await timedStore.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'stalled-preparation-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + redis.get = (() => new Promise(() => {})) as Redis['get']; + + await expect( + timedStore.dispatch({ + workerId: 'stalled-preparation-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }), + ).rejects.toThrow('Bridge worker registration read timed out'); + }); + test('keeps assignment state through deadlines longer than ten minutes', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -1012,6 +1037,53 @@ describe('RedisBridgeStore', () => { expect(await redis.exists(marker)).toBe(0); }); + test('accepts a late clean rejection and recovers its workspace fence', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'late-rejection-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const completion = store.dispatch({ + workerId: 'late-rejection-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-late-rejection', + deadlineAtMs: Date.now() + 200, + signal: new AbortController().signal, + }); + const assignment = await store.lease( + 'late-rejection-worker', + incarnationId, + 1_000, + ); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + + await store.settle( + 'late-rejection-worker', + assignment?.assignmentId ?? '', + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'rejected', + error: 'syntax_error', + }, + ); + expect( + await redis.keys( + 'codeapi:bridge:v1:worker:late-rejection-worker:workspace:*:quarantined', + ), + ).toHaveLength(0); + }); + test('releases the worker lock when generation allocation fails', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 61aa3f7a..816e1b41 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -173,6 +173,28 @@ export class RedisBridgeStore { private readonly redisCommandTimeoutMs = DEFAULT_REDIS_COMMAND_TIMEOUT_MS, ) {} + private async dispatchCommand( + command: () => Promise, + args: { deadlineAtMs: number; signal: AbortSignal }, + label: string, + ): Promise { + this.assertDispatchActive(args.signal, args.deadlineAtMs); + try { + return await boundedCommand( + command(), + Math.max( + 1, + Math.min(this.redisCommandTimeoutMs, args.deadlineAtMs - Date.now()), + ), + label, + args.signal, + ); + } catch (error) { + this.assertDispatchActive(args.signal, args.deadlineAtMs); + throw error; + } + } + async register(registration: BridgeWorkerRegistration): Promise { const script = [ 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', @@ -240,7 +262,11 @@ export class RedisBridgeStore { ) => Promise; }): Promise { this.assertDispatchActive(args.signal, args.deadlineAtMs); - let registration = await this.registration(args.workerId); + let registration = await this.dispatchCommand( + () => this.registration(args.workerId), + args, + 'Bridge worker registration read', + ); if (registration == null) { throw new BridgeStoreError( 'WORKER_OFFLINE', @@ -258,8 +284,13 @@ export class RedisBridgeStore { } if ( args.runtimeSessionId !== undefined && - (await this.redis.exists( - workspaceQuarantineKey(args.workerId, args.runtimeSessionId), + (await this.dispatchCommand( + () => + this.redis.exists( + workspaceQuarantineKey(args.workerId, args.runtimeSessionId ?? ''), + ), + args, + 'Bridge workspace fence read', )) === 1 ) { throw new BridgeStoreError( @@ -271,24 +302,33 @@ export class RedisBridgeStore { 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`, - ); - } - + const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; let resultCommitted = false; try { + const locked = await this.dispatchCommand( + () => + this.acquireLock( + args.workerId, + assignmentId, + lockIncarnationId, + ttlSeconds, + ), + args, + 'Bridge assignment lock acquisition', + ); + if (!locked) { + throw new BridgeStoreError( + 'WORKER_BUSY', + `Bridge worker ${args.workerId} is busy`, + ); + } this.assertDispatchActive(args.signal, args.deadlineAtMs); - const generation = await this.redis.incr(generationKey(args.workerId)); + const generation = await this.dispatchCommand( + () => this.redis.incr(generationKey(args.workerId)), + args, + 'Bridge assignment generation allocation', + ); assignment = { protocolVersion: BRIDGE_PROTOCOL_VERSION, assignmentId, @@ -308,9 +348,17 @@ export class RedisBridgeStore { for (let attempt = 0; attempt < 8 && !queued; attempt += 1) { this.assertDispatchActive(args.signal, args.deadlineAtMs); assignment.incarnationId = registration.incarnationId; - queued = await this.enqueueForActiveIncarnation(assignment, ttlSeconds); + queued = await this.dispatchCommand( + () => this.enqueueForActiveIncarnation(assignment!, ttlSeconds), + args, + 'Bridge assignment enqueue', + ); if (queued) break; - const replacement = await this.registration(args.workerId); + const replacement = await this.dispatchCommand( + () => this.registration(args.workerId), + args, + 'Bridge replacement registration read', + ); if (replacement == null) { throw new BridgeStoreError( 'WORKER_OFFLINE', @@ -539,7 +587,10 @@ export class RedisBridgeStore { 'Bridge assignment lease is stale', ); } - if (Date.parse(assignment.expiresAt) <= Date.now()) { + if ( + settlement.status !== 'rejected' && + Date.parse(assignment.expiresAt) <= Date.now() + ) { throw new BridgeStoreError( 'ASSIGNMENT_EXPIRED', 'Bridge assignment has expired', @@ -724,18 +775,31 @@ export class RedisBridgeStore { if (raw != null) return JSON.parse(raw) as CodeBridgeSettlement; await delay(POLL_INTERVAL_MS, signal); } + const closeKeys = [ + assignmentKey(assignment.assignmentId), + settlementKey(assignment.assignmentId), + ]; + if (assignment.runtimeSessionId !== undefined) { + closeKeys.push( + workspaceQuarantineKey( + assignment.workerId, + assignment.runtimeSessionId, + ), + ); + } const closeScript = [ 'local settlement = redis.call(\'GET\', KEYS[2])', 'if settlement then return settlement end', + 'if #KEYS == 3 and redis.call(\'GET\', KEYS[3]) == ARGV[1] then return nil end', 'redis.call(\'DEL\', KEYS[1])', 'return nil', ].join('\n'); const finalSettlement = await boundedCommand( this.redis.eval( closeScript, - 2, - assignmentKey(assignment.assignmentId), - settlementKey(assignment.assignmentId), + closeKeys.length, + ...closeKeys, + assignment.assignmentId, ), this.redisCommandTimeoutMs, 'Bridge settlement close', @@ -833,16 +897,16 @@ export class RedisBridgeStore { assignmentId: string, assignment: StoredAssignment | undefined, ): Promise { - await this.cancel(assignmentId); - if (assignment == null) { - await boundedCommand( - this.releaseLock(workerId, assignmentId), - this.redisCommandTimeoutMs, - 'Bridge assignment lock release', - ); - return; - } - await this.cleanup(assignment); + await Promise.all([ + this.cancel(assignmentId), + assignment == null + ? boundedCommand( + this.releaseLock(workerId, assignmentId), + this.redisCommandTimeoutMs, + 'Bridge assignment lock release', + ) + : this.cleanup(assignment), + ]); } private async commitPendingWorkspace( @@ -916,6 +980,9 @@ export class RedisBridgeStore { 'if queued > 0 and #KEYS == 3 and redis.call(\'GET\', KEYS[3]) == ARGV[1] then', " redis.call('DEL', KEYS[3])", 'end', + 'if queued == 0 and #KEYS == 3 and redis.call(\'GET\', KEYS[3]) == ARGV[1] then', + ' return 0', + 'end', "return redis.call('DEL', KEYS[1])", ].join('\n'); await Promise.all([ diff --git a/service/src/local-api.ts b/service/src/local-api.ts index 701270d7..21c607d3 100644 --- a/service/src/local-api.ts +++ b/service/src/local-api.ts @@ -10,6 +10,7 @@ import express, { json, Router } from 'express'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; +import bridgeRouter from './bridge/router'; import { requestErrorLogger, requestNotFoundLogger } from './middleware/request-error-logger'; import { executionProfileMiddleware } from './middleware/execution-profile'; import { localAuth } from './auth/local'; @@ -45,6 +46,7 @@ app.get('/v1/health', async (_, res) => { } }); +v1.use('/bridge', bridgeRouter); v1.use(localAuth); v1.use(serviceRouter); v1.use(programmaticRouter); From 159be6aad852823c67b37522f805eea8ae14cc03 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 12:07:52 -0400 Subject: [PATCH 17/29] fix: acknowledge bridge lease delivery --- docs/remote-bridge/README.md | 3 + packages/code/src/worker.test.ts | 45 ++++++++- packages/code/src/worker.ts | 59 +++++++++++- service/src/bridge/router.ts | 57 ++++++++---- service/src/bridge/store.test.ts | 77 ++++++++++++++-- service/src/bridge/store.ts | 151 +++++++++++++++++++++++++++---- service/src/local-api.ts | 8 +- 7 files changed, 351 insertions(+), 49 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index b67fc29b..dd22e73b 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -88,6 +88,9 @@ execution. the currently registered incarnation. - Request cancellation is polled by the worker and aborts the local sandbox request. +- A leased assignment remains in a Redis-backed delivery claim until the worker + explicitly acknowledges it; reconnecting before acknowledgement redelivers + the same fenced assignment instead of losing it after an HTTP disconnect. - The sandbox receives the stable runtime session ID separately from the lease; workspace state belongs to that session, not to a transient assignment. diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 925be950..e308b571 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -108,10 +108,39 @@ test('worker acknowledges a discarded workspace through the reset endpoint', asy }); }); +test('worker bounds a stalled workspace reset request', async () => { + 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: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + resetTransportTimeoutMs: 20, + fetchImpl: async (_input, init) => + await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }), + }); + + await assert.rejects(worker.resetWorkspace('rt-user-1'), { + name: 'AbortError', + }); +}); + test('worker continues after an assignment-scoped settlement conflict', async () => { const controller = new AbortController(); let registrations = 0; let leases = 0; + let leaseAcknowledged = false; let observedError: unknown; const assignment: BridgeAssignment = { protocolVersion: 1, @@ -150,7 +179,15 @@ test('worker continues after an assignment-scoped settlement conflict', async () { status: 200, headers: { 'Content-Type': 'application/json' } }, ); } + if (url.endsWith('/ack')) { + leaseAcknowledged = true; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } if (url.endsWith('/execute')) { + assert.equal(leaseAcknowledged, true); return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { status: 200, headers: { 'Content-Type': 'application/json' }, @@ -948,7 +985,13 @@ test('worker subtracts lease response transit from the server budget', async () sandboxProfile: 'nsjail', runtimes: ['bash'], }, - fetchImpl: async () => { + fetchImpl: async (input) => { + if (String(input).endsWith('/ack')) { + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } now += 50; return new Response( JSON.stringify({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index a47fadb6..40382d81 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -24,6 +24,8 @@ export interface BridgeWorkerOptions { leaseWaitMs?: number; leaseTransportGraceMs?: number; registrationTransportTimeoutMs?: number; + leaseAckTransportTimeoutMs?: number; + resetTransportTimeoutMs?: number; cancellationPollIntervalMs?: number; cancellationTransportTimeoutMs?: number; rejectionAckGraceMs?: number; @@ -39,6 +41,7 @@ const DEFAULT_LEASE_TRANSPORT_GRACE_MS = 5_000; const DEFAULT_RECONNECT_DELAY_MS = 1_000; const DEFAULT_REGISTRATION_TTL_MS = 60_000; const DEFAULT_REGISTRATION_TRANSPORT_TIMEOUT_MS = 10_000; +const DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS = 10_000; const DEFAULT_CANCELLATION_POLL_INTERVAL_MS = 500; const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; @@ -138,7 +141,7 @@ export class BridgeWorker { if (runtimeSessionId.trim().length === 0) { throw new BridgeProtocolError('Runtime session ID is required'); } - await this.request( + await this.timedRequest( `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -146,6 +149,11 @@ export class BridgeWorker { runtimeSessionId, confirmDiscarded: true, }, + Math.max( + 1, + this.options.resetTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + ), signal, ); } @@ -217,13 +225,37 @@ export class BridgeWorker { 0, Date.now() - requestStartedAtMs - (response.serverElapsedMs ?? 0), ); - return { + const adjustedAssignment = { ...response.assignment, remainingMs: Math.max( 0, (response.assignment.remainingMs ?? 0) - transportElapsedMs, ), }; + const acknowledgementStartedAtMs = Date.now(); + await this.timedRequest( + this.assignmentUrl(adjustedAssignment, 'ack'), + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + generation: adjustedAssignment.generation, + leaseToken: adjustedAssignment.leaseToken, + }, + Math.max( + 1, + this.options.leaseAckTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + ), + signal, + ); + return { + ...adjustedAssignment, + remainingMs: Math.max( + 0, + (adjustedAssignment.remainingMs ?? 0) - + (Date.now() - acknowledgementStartedAtMs), + ), + }; } async run(signal?: AbortSignal): Promise { @@ -630,4 +662,27 @@ export class BridgeWorker { } return payload as T; } + + private async timedRequest( + url: string, + body: object, + timeoutMs: number, + signal?: AbortSignal, + ): Promise { + const controller = new AbortController(); + const abortRequest = (): void => controller.abort(); + if (signal?.aborted) { + abortRequest(); + } else { + signal?.addEventListener('abort', abortRequest, { once: true }); + } + const timeout = setTimeout(abortRequest, timeoutMs); + timeout.unref?.(); + try { + return await this.request(url, body, controller.signal); + } finally { + clearTimeout(timeout); + signal?.removeEventListener('abort', abortRequest); + } + } } diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 88142443..2d25a1b0 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -240,26 +240,11 @@ router.post( if (assignment != null) await bridgeStore.returnLease(assignment); return; } - const delivered = await new Promise((resolve) => { - const onFinish = (): void => { - res.off('close', onClose); - resolve(true); - }; - const onClose = (): void => { - res.off('finish', onFinish); - resolve(false); - }; - res.once('finish', onFinish); - res.once('close', onClose); - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), - assignment, - }); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), + assignment, }); - if (!delivered && assignment != null) { - await bridgeStore.returnLease(assignment); - } } finally { req.off('aborted', abortLease); res.off('close', abortLease); @@ -274,6 +259,40 @@ router.post( }), ); +router.post( + '/workers/:workerId/assignments/:assignmentId/ack', + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isSafeInteger(body.generation) || + Number(body.generation) < 1 || + typeof body.leaseToken !== 'string' || + body.leaseToken.length < 32 + ) { + res.status(400).json({ error: 'Invalid bridge lease acknowledgement' }); + return; + } + try { + await bridgeStore.acknowledgeLease( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + Number(body.generation), + body.leaseToken, + ); + 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/settle', asyncRoute(async (req, res) => { diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index f637ace4..bd10a8d2 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -69,6 +69,50 @@ describe('RedisBridgeStore', () => { }); }); + test('redelivers a lease claim until the worker acknowledges it', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'claim-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const completion = store.dispatch({ + workerId: 'claim-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + const first = await store.lease('claim-worker', incarnationId, 1_000); + const redelivered = await store.lease( + 'claim-worker', + incarnationId, + 1_000, + ); + expect(redelivered?.assignmentId).toBe(first?.assignmentId); + + await store.acknowledgeLease( + 'claim-worker', + incarnationId, + first?.assignmentId ?? '', + first?.generation ?? 0, + first?.leaseToken ?? '', + ); + await store.settle('claim-worker', first?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: first?.generation ?? 0, + leaseToken: first?.leaseToken ?? '', + incarnationId, + status: 'rejected', + error: 'test complete', + }); + await expect(completion).resolves.toMatchObject({ status: 'rejected' }); + }); + test('rejects dispatch to an offline worker', async () => { const controller = new AbortController(); await expect( @@ -255,16 +299,23 @@ describe('RedisBridgeStore', () => { signal: dispatchController.signal, }); const leaseController = new AbortController(); - redis.lpop = (async (...args: Parameters) => { - const assignmentId = await redisLpop(...args); - if (assignmentId != null) leaseController.abort(); - return assignmentId; - }) as Redis['lpop']; + redis.eval = (async (...args: Parameters) => { + const result = await redisEval(...args); + if ( + String(args[0]).includes( + "local claimed = redis.call('GET', KEYS[2])", + ) && + result != null + ) { + leaseController.abort(); + } + return result; + }) as Redis['eval']; await expect( store.lease('vm-1', incarnationId, 1_000, leaseController.signal), ).resolves.toBeUndefined(); - redis.lpop = redisLpop as Redis['lpop']; + redis.eval = redisEval as Redis['eval']; const recovered = await store.lease('vm-1', incarnationId, 1_000); expect(recovered).toBeDefined(); @@ -932,6 +983,13 @@ describe('RedisBridgeStore', () => { 1_000, ); expect(assignment).toBeDefined(); + await store.acknowledgeLease( + 'lost-worker', + incarnationId, + assignment?.assignmentId ?? '', + assignment?.generation ?? 0, + assignment?.leaseToken ?? '', + ); const [marker] = await redis.keys( 'codeapi:bridge:v1:worker:lost-worker:workspace:*:quarantined', ); @@ -1061,6 +1119,13 @@ describe('RedisBridgeStore', () => { incarnationId, 1_000, ); + await store.acknowledgeLease( + 'late-rejection-worker', + incarnationId, + assignment?.assignmentId ?? '', + assignment?.generation ?? 0, + assignment?.leaseToken ?? '', + ); await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED', }); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 816e1b41..eedec8d1 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -77,6 +77,14 @@ function queueKey(workerId: string, incarnationId: string): string { return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:assignments`; } +function leaseClaimKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:lease-claim`; +} + +function leaseAckKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:lease-ack`; +} + function generationKey(workerId: string): string { return `${PREFIX}:worker:${workerId}:generation`; } @@ -428,9 +436,7 @@ export class RedisBridgeStore { ): Promise { const deadline = Date.now() + waitMs; while (!signalAborted(signal) && Date.now() < deadline) { - const assignmentId = await this.redis.lpop( - queueKey(workerId, incarnationId), - ); + const assignmentId = await this.claimOrPopLease(workerId, incarnationId); if (assignmentId == null) { await delay( Math.min(POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), @@ -440,8 +446,14 @@ export class RedisBridgeStore { } try { const assignment = await this.readAssignment(assignmentId); - if (assignment == null || assignment.workerId !== workerId) continue; - if (assignment.incarnationId !== incarnationId) continue; + if ( + assignment == null || + assignment.workerId !== workerId || + assignment.incarnationId !== incarnationId + ) { + await this.discardLeaseClaim(workerId, incarnationId, assignmentId); + continue; + } if (signalAborted(signal)) { await this.returnLease(assignment); return undefined; @@ -455,6 +467,7 @@ export class RedisBridgeStore { } if (Date.parse(assignment.expiresAt) <= Date.now()) { await this.clearUndeliveredWorkspaceFence(assignment); + await this.discardLeaseClaim(workerId, incarnationId, assignmentId); continue; } if (signalAborted(signal)) { @@ -481,6 +494,91 @@ export class RedisBridgeStore { return undefined; } + async acknowledgeLease( + workerId: string, + incarnationId: string, + assignmentId: string, + generation: number, + leaseToken: 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 || + assignment.generation !== generation || + tokenHash(leaseToken) !== assignment.leaseTokenHash + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment lease acknowledgement is stale', + ); + } + const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); + const acknowledged = Number( + await this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", + "redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2])", + 'return 1', + ].join('\n'), + 2, + leaseClaimKey(workerId, incarnationId), + leaseAckKey(workerId, incarnationId), + assignmentId, + String(ttlSeconds), + ), + ); + if (acknowledged !== 1) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment is not the active lease claim', + ); + } + } + + private async claimOrPopLease( + workerId: string, + incarnationId: string, + ): Promise { + const result = await this.redis.eval( + [ + "local claimed = redis.call('GET', KEYS[2])", + 'if claimed then return claimed end', + "local ttl = redis.call('TTL', KEYS[1])", + "local assignment = redis.call('LPOP', KEYS[1])", + 'if not assignment then return nil end', + "redis.call('SET', KEYS[2], assignment, 'EX', math.max(1, ttl))", + 'return assignment', + ].join('\n'), + 2, + queueKey(workerId, incarnationId), + leaseClaimKey(workerId, incarnationId), + ); + return result == null ? null : String(result); + } + + private async discardLeaseClaim( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { + await this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) == ARGV[1] then", + " return redis.call('DEL', KEYS[1], KEYS[2])", + 'end', + 'return 0', + ].join('\n'), + 2, + leaseClaimKey(workerId, incarnationId), + leaseAckKey(workerId, incarnationId), + assignmentId, + ); + } + async returnLease(assignment: CodeBridgeAssignment): Promise { await this.returnLeaseById( assignment.workerId, @@ -497,13 +595,17 @@ export class RedisBridgeStore { await this.redis.eval( [ "if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end", + "if redis.call('GET', KEYS[3]) ~= ARGV[1] then return 0 end", + "redis.call('DEL', KEYS[3], KEYS[4])", "redis.call('LREM', KEYS[2], 0, ARGV[1])", "redis.call('LPUSH', KEYS[2], ARGV[1])", 'return 1', ].join('\n'), - 2, + 4, assignmentKey(assignmentId), queueKey(workerId, incarnationId), + leaseClaimKey(workerId, incarnationId), + leaseAckKey(workerId, incarnationId), assignmentId, ); } @@ -600,6 +702,8 @@ export class RedisBridgeStore { const settlementKeys = [ assignmentKey(assignmentId), settlementKey(assignmentId), + leaseClaimKey(workerId, assignment.incarnationId), + leaseAckKey(workerId, assignment.incarnationId), ]; if (assignment.runtimeSessionId !== undefined) { settlementKeys.push( @@ -613,9 +717,10 @@ export class RedisBridgeStore { ' return -1', 'end', 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return 0 end', - 'if #KEYS == 3 and redis.call(\'GET\', KEYS[3]) ~= ARGV[3] then return -2 end', + 'if #KEYS == 5 and redis.call(\'GET\', KEYS[5]) ~= ARGV[3] then return -2 end', 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', - 'if #KEYS == 3 and ARGV[4] == \"rejected\" then redis.call(\'DEL\', KEYS[3]) end', + 'if redis.call(\'GET\', KEYS[3]) == ARGV[3] then redis.call(\'DEL\', KEYS[3], KEYS[4]) end', + 'if #KEYS == 5 and ARGV[4] == \"rejected\" then redis.call(\'DEL\', KEYS[5]) end', 'return 1', ].join('\n'); const accepted = Number( @@ -966,24 +1071,29 @@ export class RedisBridgeStore { const keys = [ assignmentKey(assignment.assignmentId), queueKey(assignment.workerId, assignment.incarnationId), + leaseClaimKey(assignment.workerId, assignment.incarnationId), + leaseAckKey(assignment.workerId, assignment.incarnationId), + assignment.runtimeSessionId === undefined + ? `${assignmentKey(assignment.assignmentId)}:no-workspace` + : workspaceQuarantineKey( + assignment.workerId, + assignment.runtimeSessionId, + ), ]; - if (assignment.runtimeSessionId !== undefined) { - keys.push( - workspaceQuarantineKey( - assignment.workerId, - assignment.runtimeSessionId, - ), - ); - } const cleanupScript = [ "local queued = redis.call('LREM', KEYS[2], 0, ARGV[1])", - 'if queued > 0 and #KEYS == 3 and redis.call(\'GET\', KEYS[3]) == ARGV[1] then', - " redis.call('DEL', KEYS[3])", + "local claimed = redis.call('GET', KEYS[3]) == ARGV[1]", + "local acknowledged = redis.call('GET', KEYS[4]) == ARGV[1]", + 'if ARGV[2] == "1" and (queued > 0 or (claimed and not acknowledged)) and redis.call(\'GET\', KEYS[5]) == ARGV[1] then', + " redis.call('DEL', KEYS[5])", + 'end', + 'if claimed and not acknowledged then', + " redis.call('DEL', KEYS[3], KEYS[4])", 'end', - 'if queued == 0 and #KEYS == 3 and redis.call(\'GET\', KEYS[3]) == ARGV[1] then', + 'if queued == 0 and acknowledged and ARGV[2] == "1" and redis.call(\'GET\', KEYS[5]) == ARGV[1] then', ' return 0', 'end', - "return redis.call('DEL', KEYS[1])", + "return redis.call('DEL', KEYS[1], KEYS[3], KEYS[4])", ].join('\n'); await Promise.all([ boundedCommand( @@ -992,6 +1102,7 @@ export class RedisBridgeStore { keys.length, ...keys, assignment.assignmentId, + assignment.runtimeSessionId === undefined ? '0' : '1', ), this.redisCommandTimeoutMs, 'Bridge assignment cleanup', diff --git a/service/src/local-api.ts b/service/src/local-api.ts index 21c607d3..aafa9f03 100644 --- a/service/src/local-api.ts +++ b/service/src/local-api.ts @@ -21,7 +21,11 @@ import './workers'; import { env } from './config'; import logger from './logger'; import { shutdownTelemetry, traceHttpRequest } from './telemetry'; -import { validateExecutionProfilePolicy } from './secure-startup'; +import { + validateApiBridgePolicy, + validateExecutionProfilePolicy, + validateSandboxBackendPolicy, +} from './secure-startup'; import { configureExecutionProfileMetrics } from './metrics'; const app = express(); @@ -58,7 +62,9 @@ app.use(requestErrorLogger); async function localStartup(): Promise { logger.info('Starting local development server...'); logger.info('⚠️ LOCAL MODE - No authentication required'); + validateApiBridgePolicy(); validateExecutionProfilePolicy(); + validateSandboxBackendPolicy(); configureExecutionProfileMetrics({ profile: env.EXECUTION_PROFILE, sandboxBackend: env.SANDBOX_BACKEND, From 5876b83aa42dc67fcca4fd96f5051cc4e2ea6077 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 12:23:22 -0400 Subject: [PATCH 18/29] fix: close bridge deadline gaps --- packages/code/src/worker.test.ts | 96 ++++++++++++++++++++++++++++++++ packages/code/src/worker.ts | 22 ++++++-- service/src/bridge/store.test.ts | 79 ++++++++++++++++++++++++++ service/src/bridge/store.ts | 32 ++++++++--- 4 files changed, 215 insertions(+), 14 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index e308b571..5e3bcfd8 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -1021,6 +1021,102 @@ test('worker subtracts lease response transit from the server budget', async () } }); +test('worker rejects a lease whose acknowledgement exhausts its budget', async () => { + const originalNow = Date.now; + let now = 20_000; + Date.now = () => now; + try { + 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: async (input) => { + if (String(input).endsWith('/ack')) { + now += 10; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + return new Response( + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 0, + assignment: { + protocolVersion: 1, + assignmentId: 'expired-after-ack', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(0).toISOString(), + remainingMs: 10, + request: { body: { language: 'bash' }, headers: {} }, + }, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await assert.rejects(worker.lease(), /expired during lease acknowledgement/); + } finally { + Date.now = originalNow; + } +}); + +test('worker clamps rejected settlement errors to the protocol limit', async () => { + let rejection = ''; + 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: async (input, init) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ error: 'x'.repeat(5_000) }), { + status: 400, + headers: { 'Content-Type': 'application/json' }, + }); + } + const settlement = JSON.parse(String(init?.body)) as { error: string }; + rejection = settlement.error; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'long-rejection', + 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(), + remainingMs: 1_000, + runtimeSessionId: 'rt-long-rejection', + request: { body: { language: 'bash' }, headers: {} }, + }); + assert.equal(rejection.length, 4_096); +}); + test('worker quarantines an explicitly dirty stateful sandbox response', async () => { let settlementAttempted = false; const worker = new BridgeWorker({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 40382d81..494e5cd9 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -47,6 +47,7 @@ const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; const SETTLEMENT_RETRY_DELAY_MS = 100; const REJECTION_ACK_GRACE_MS = 30_000; +const MAX_SETTLEMENT_ERROR_LENGTH = 4_096; const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; function normalizedBaseUrl(value: string): string { @@ -248,13 +249,19 @@ export class BridgeWorker { ), signal, ); + const remainingMs = Math.max( + 0, + (adjustedAssignment.remainingMs ?? 0) - + (Date.now() - acknowledgementStartedAtMs), + ); + if (remainingMs <= 0) { + throw new BridgeProtocolError( + 'Bridge assignment expired during lease acknowledgement', + ); + } return { ...adjustedAssignment, - remainingMs: Math.max( - 0, - (adjustedAssignment.remainingMs ?? 0) - - (Date.now() - acknowledgementStartedAtMs), - ), + remainingMs, }; } @@ -379,7 +386,10 @@ export class BridgeWorker { incarnationId: this.incarnationId, status: 'rejected', error: - error instanceof Error ? error.message : 'Sandbox execution failed', + (error instanceof Error + ? error.message + : 'Sandbox execution failed' + ).slice(0, MAX_SETTLEMENT_ERROR_LENGTH), }; } diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index bd10a8d2..8e14388e 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -957,6 +957,85 @@ describe('RedisBridgeStore', () => { expect(await redis.exists(pendingMarker)).toBe(0); }); + test('bounds a stalled Redis workspace commit command', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 20); + await timedStore.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'stalled-commit-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + let releaseFinalizer!: () => void; + const finalizerGate = new Promise((resolve) => { + releaseFinalizer = resolve; + }); + let finalizerStarted!: () => void; + const started = new Promise((resolve) => { + finalizerStarted = resolve; + }); + const completion = timedStore.dispatch({ + workerId: 'stalled-commit-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-stalled-commit', + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + finalize: async (settlement) => { + finalizerStarted(); + await finalizerGate; + return settlement; + }, + }); + const assignment = await timedStore.lease( + 'stalled-commit-worker', + incarnationId, + 1_000, + ); + await timedStore.acknowledgeLease( + 'stalled-commit-worker', + incarnationId, + assignment?.assignmentId ?? '', + assignment?.generation ?? 0, + assignment?.leaseToken ?? '', + ); + await timedStore.settle( + 'stalled-commit-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-stalled-commit', + files: [], + }, + }, + ); + await started; + redis.eval = ((...args: Parameters) => { + if ( + Number(args[1]) === 1 && + String(args[0]).includes("return redis.call('DEL', KEYS[1])") + ) { + return new Promise(() => {}); + } + return redisEval(...args); + }) as Redis['eval']; + releaseFinalizer(); + + await expect(completion).rejects.toThrow( + 'Bridge workspace commit timed out', + ); + }); + test('keeps an in-flight workspace fenced when execution never settles', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index eedec8d1..b0f36112 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -400,7 +400,12 @@ export class RedisBridgeStore { args.finalize == null ? settlement : await args.finalize(settlement); - await this.commitPendingWorkspace(assignment, settlement); + await this.commitPendingWorkspace( + assignment, + settlement, + args.deadlineAtMs, + args.signal, + ); resultCommitted = true; return result; } catch (error) { @@ -1017,6 +1022,8 @@ export class RedisBridgeStore { private async commitPendingWorkspace( assignment: StoredAssignment, settlement: CodeBridgeSettlement, + deadlineAtMs: number, + signal: AbortSignal, ): Promise { if ( assignment.runtimeSessionId === undefined || @@ -1024,6 +1031,7 @@ export class RedisBridgeStore { ) { return; } + const runtimeSessionId = assignment.runtimeSessionId; const script = [ 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', ' return redis.call(\'DEL\', KEYS[1])', @@ -1031,14 +1039,22 @@ export class RedisBridgeStore { 'return 0', ].join('\n'); const committed = Number( - await this.redis.eval( - script, - 1, - workspaceQuarantineKey( - assignment.workerId, - assignment.runtimeSessionId, + await boundedCommand( + this.redis.eval( + script, + 1, + workspaceQuarantineKey( + assignment.workerId, + runtimeSessionId, + ), + assignment.assignmentId, ), - assignment.assignmentId, + Math.max( + 1, + Math.min(this.redisCommandTimeoutMs, deadlineAtMs - Date.now()), + ), + 'Bridge workspace commit', + signal, ), ); if (committed !== 1) { From 605f66a31108835647e5ce1bfd6581c8f654eb83 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 12:44:31 -0400 Subject: [PATCH 19/29] fix: harden bridge recovery edges --- packages/code/src/worker.test.ts | 76 +++++++++++++++++++++++++++++++- packages/code/src/worker.ts | 26 +++++++++-- service/src/bridge/store.test.ts | 57 ++++++++++++++++++++++++ service/src/bridge/store.ts | 21 ++++++--- 4 files changed, 169 insertions(+), 11 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 5e3bcfd8..468ba2e1 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -682,7 +682,21 @@ test('worker keeps a definite stateful rejection nonfatal when settlement is amb test('worker retries a known-clean rejection after shutdown until acknowledged', async () => { const controller = new AbortController(); let settlementAttempts = 0; + let registrations = 0; const fetchImpl: typeof fetch = async (input) => { + if (String(input).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 (String(input).endsWith('/execute')) { return new Response(JSON.stringify({ error: 'syntax_error' }), { status: 400, @@ -715,6 +729,7 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', fetchImpl, }); + await worker.register(); await worker.executeAndSettle( { protocolVersion: 1, @@ -732,6 +747,54 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', ); assert.equal(controller.signal.aborted, true); assert.equal(settlementAttempts, 2); + assert.ok(registrations > 1); +}); + +test('worker quarantines a stateful workspace after a sandbox 5xx response', async () => { + let settlementAttempted = false; + 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: async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ error: 'upstream failed' }), { + status: 502, + headers: { 'Content-Type': 'application/json' }, + }); + } + settlementAttempted = true; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'ambiguous-5xx', + 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(), + remainingMs: 1_000, + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); + assert.equal(settlementAttempted, false); }); test('worker quarantines a stateful workspace after the sandbox request aborts', async () => { @@ -1024,6 +1087,7 @@ test('worker subtracts lease response transit from the server budget', async () test('worker rejects a lease whose acknowledgement exhausts its budget', async () => { const originalNow = Date.now; let now = 20_000; + let abandonedSettlement: Record | undefined; Date.now = () => now; try { const worker = new BridgeWorker({ @@ -1037,7 +1101,7 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( sandboxProfile: 'nsjail', runtimes: ['bash'], }, - fetchImpl: async (input) => { + fetchImpl: async (input, init) => { if (String(input).endsWith('/ack')) { now += 10; return new Response( @@ -1045,6 +1109,15 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( { status: 200, headers: { 'Content-Type': 'application/json' } }, ); } + if (String(input).endsWith('/settle')) { + abandonedSettlement = JSON.parse( + String(init?.body), + ) as Record; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } return new Response( JSON.stringify({ protocolVersion: 1, @@ -1067,6 +1140,7 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( }); await assert.rejects(worker.lease(), /expired during lease acknowledgement/); + assert.equal(abandonedSettlement?.status, 'rejected'); } finally { Date.now = originalNow; } diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 494e5cd9..28451f79 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -255,6 +255,22 @@ export class BridgeWorker { (Date.now() - acknowledgementStartedAtMs), ); if (remainingMs <= 0) { + await this.settleWithRetry( + adjustedAssignment, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: adjustedAssignment.generation, + leaseToken: adjustedAssignment.leaseToken, + incarnationId: this.incarnationId, + status: 'rejected', + error: 'Bridge assignment expired during lease acknowledgement', + }, + Date.now() + + Math.max( + 0, + this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, + ), + ); throw new BridgeProtocolError( 'Bridge assignment expired during lease acknowledgement', ); @@ -319,7 +335,6 @@ export class BridgeWorker { let heartbeatError: unknown; const heartbeat = this.maintainRegistration( heartbeatController.signal, - executionController, ).catch((error) => { heartbeatError = error; executionController.abort(); @@ -357,6 +372,10 @@ export class BridgeWorker { if (heartbeatError != null) throw heartbeatError; if (!response.ok) { sandboxRejectedExecution = + response.status >= 400 && + response.status < 500 && + response.status !== 408 && + response.status !== 429 && errorMessage(payload) !== 'session_workspace_dirty'; throw new BridgeProtocolError( errorMessage(payload) ?? @@ -474,9 +493,8 @@ export class BridgeWorker { private async maintainRegistration( signal: AbortSignal, - executionController: AbortController, ): Promise { - while (!signal.aborted && !executionController.signal.aborted) { + while (!signal.aborted) { const heartbeatIntervalMs = Math.max( MIN_REGISTRATION_HEARTBEAT_MS, Math.floor(this.registrationTtlMs / 2), @@ -488,7 +506,7 @@ export class BridgeWorker { ), signal, ); - if (signal.aborted || executionController.signal.aborted) return; + if (signal.aborted) return; await this.register(signal); } } diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 8e14388e..afa44526 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -113,6 +113,63 @@ describe('RedisBridgeStore', () => { await expect(completion).resolves.toMatchObject({ status: 'rejected' }); }); + test('performs one immediate lease poll when wait is zero', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'nonblocking-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const completion = store.dispatch({ + workerId: 'nonblocking-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + for (let attempt = 0; attempt < 100; attempt += 1) { + if ( + ( + await redis.keys( + 'codeapi:bridge:v1:assignment:*', + ) + ).length > 0 + ) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 1)); + } + + const assignment = await store.lease( + 'nonblocking-worker', + incarnationId, + 0, + ); + expect(assignment).toBeDefined(); + await store.settle('nonblocking-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'rejected', + error: 'test complete', + }); + await expect(completion).resolves.toMatchObject({ status: 'rejected' }); + }); + + test('bounds a stalled quarantine command', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.eval = (() => new Promise(() => undefined)) as Redis['eval']; + + await expect( + timedStore.quarantine('stalled-worker', incarnationId, 'rt-user-1'), + ).rejects.toThrow('Bridge worker quarantine timed out'); + }); + test('rejects dispatch to an offline worker', async () => { const controller = new AbortController(); await expect( diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index b0f36112..c30f74c8 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -440,7 +440,12 @@ export class RedisBridgeStore { signal?: AbortSignal, ): Promise { const deadline = Date.now() + waitMs; - while (!signalAborted(signal) && Date.now() < deadline) { + let firstPoll = true; + while ( + !signalAborted(signal) && + (firstPoll || Date.now() < deadline) + ) { + firstPoll = false; const assignmentId = await this.claimOrPopLease(workerId, incarnationId); if (assignmentId == null) { await delay( @@ -799,11 +804,15 @@ export class RedisBridgeStore { if (runtimeSessionId !== undefined) { keys.push(workspaceQuarantineKey(workerId, runtimeSessionId)); } - await this.redis.eval( - script, - keys.length, - ...keys, - incarnationId, + await boundedCommand( + this.redis.eval( + script, + keys.length, + ...keys, + incarnationId, + ), + this.redisCommandTimeoutMs, + 'Bridge worker quarantine', ); } From d5722b0ff20fe35731086e87b2345ac64385df63 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 13:14:46 -0400 Subject: [PATCH 20/29] fix: preserve bridge rejection recovery --- packages/code/src/worker.test.ts | 90 +++++++++++++++++++++++++++++++- packages/code/src/worker.ts | 43 +++++++++------ service/src/utils.test.ts | 14 +++++ service/src/utils.ts | 2 +- 4 files changed, 130 insertions(+), 19 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 468ba2e1..17451aa8 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -750,6 +750,70 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', assert.ok(registrations > 1); }); +test('worker preserves a definite rejection when its heartbeat fails', async () => { + let registrations = 0; + let rejectedSettlement = false; + 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: async (input, init) => { + if (String(input).endsWith('/workers/register')) { + registrations += 1; + if (registrations > 1) throw new TypeError('registration unavailable'); + 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 (String(input).endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 40)); + return new Response(JSON.stringify({ error: 'syntax_error' }), { + status: 400, + headers: { 'Content-Type': 'application/json' }, + }); + } + rejectedSettlement = + JSON.parse(String(init?.body) || '{}').status === 'rejected'; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await worker.register(); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'clean-rejection-after-heartbeat-error', + 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(), + remainingMs: 1_000, + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.ok(registrations > 1); + assert.equal(rejectedSettlement, true); +}); + test('worker quarantines a stateful workspace after a sandbox 5xx response', async () => { let settlementAttempted = false; const worker = new BridgeWorker({ @@ -1086,8 +1150,10 @@ test('worker subtracts lease response transit from the server budget', async () test('worker rejects a lease whose acknowledgement exhausts its budget', async () => { const originalNow = Date.now; - let now = 20_000; + let now = 100_000; let abandonedSettlement: Record | undefined; + let registrations = 0; + let settlementAttempts = 0; Date.now = () => now; try { const worker = new BridgeWorker({ @@ -1102,6 +1168,19 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( runtimes: ['bash'], }, fetchImpl: async (input, init) => { + if (String(input).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 (String(input).endsWith('/ack')) { now += 10; return new Response( @@ -1110,9 +1189,16 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( ); } if (String(input).endsWith('/settle')) { + settlementAttempts += 1; abandonedSettlement = JSON.parse( String(init?.body), ) as Record; + if (settlementAttempts === 1) { + return new Response(JSON.stringify({ error: 'unavailable' }), { + status: 503, + headers: { 'Content-Type': 'application/json' }, + }); + } return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), { status: 200, headers: { 'Content-Type': 'application/json' } }, @@ -1141,6 +1227,8 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( await assert.rejects(worker.lease(), /expired during lease acknowledgement/); assert.equal(abandonedSettlement?.status, 'rejected'); + assert.ok(registrations > 0); + assert.equal(settlementAttempts, 2); } finally { Date.now = originalNow; } diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 28451f79..b86afa68 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -255,22 +255,31 @@ export class BridgeWorker { (Date.now() - acknowledgementStartedAtMs), ); if (remainingMs <= 0) { - await this.settleWithRetry( - adjustedAssignment, - { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - generation: adjustedAssignment.generation, - leaseToken: adjustedAssignment.leaseToken, - incarnationId: this.incarnationId, - status: 'rejected', - error: 'Bridge assignment expired during lease acknowledgement', - }, - Date.now() + - Math.max( - 0, - this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, - ), - ); + const heartbeatController = new AbortController(); + const heartbeat = this.maintainRegistration( + heartbeatController.signal, + ).catch(() => undefined); + try { + await this.settleWithRetry( + adjustedAssignment, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: adjustedAssignment.generation, + leaseToken: adjustedAssignment.leaseToken, + incarnationId: this.incarnationId, + status: 'rejected', + error: 'Bridge assignment expired during lease acknowledgement', + }, + Date.now() + + Math.max( + 0, + this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, + ), + ); + } finally { + heartbeatController.abort(); + await heartbeat; + } throw new BridgeProtocolError( 'Bridge assignment expired during lease acknowledgement', ); @@ -369,7 +378,6 @@ export class BridgeWorker { }, ); const payload = (await response.json()) as object; - if (heartbeatError != null) throw heartbeatError; if (!response.ok) { sandboxRejectedExecution = response.status >= 400 && @@ -383,6 +391,7 @@ export class BridgeWorker { response.status, ); } + if (heartbeatError != null) throw heartbeatError; settlement = { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, diff --git a/service/src/utils.test.ts b/service/src/utils.test.ts index 8aa0dac0..879736c5 100644 --- a/service/src/utils.test.ts +++ b/service/src/utils.test.ts @@ -169,6 +169,20 @@ describe('sandbox error formatting', () => { } }); + test('maps multiline remote bridge failures without exposing details', () => { + const failure = publicExecutionFailure( + new Error('BRIDGE_EXECUTION_FAILED: first line\nprivate second line'), + ); + expect(failure).toEqual({ + status: 502, + body: { + error: 'bridge_execution_failed', + message: 'Remote code execution failed', + }, + }); + expect(JSON.stringify(failure)).not.toContain('private second line'); + }); + 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 b94d8fd2..e47a7d4c 100644 --- a/service/src/utils.ts +++ b/service/src/utils.ts @@ -131,7 +131,7 @@ export function publicExecutionFailure(error: unknown): { status: number; body: * 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_]+|BRIDGE_[A-Z_]+|SESSION_INPUT_[A-Z_]+):\s*(.+)$/, + /^(RUNTIME_SESSION_BUSY|MICROVM_[A-Z_]+|BRIDGE_[A-Z_]+|SESSION_INPUT_[A-Z_]+):/, ); if (backendMatch) { const code = backendMatch[1]; From ca68fe6250eaefa6516ae7fc9c5fd72de31c3584 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 13:37:02 -0400 Subject: [PATCH 21/29] fix: fence bridge lease recovery --- packages/code/src/protocol.ts | 5 + packages/code/src/worker.test.ts | 82 ++++++++++++++- packages/code/src/worker.ts | 160 ++++++++++++++++++++--------- service/src/bridge/router.ts | 8 +- service/src/bridge/store.test.ts | 102 ++++++++++++++++++ service/src/bridge/store.ts | 100 ++++++++++++++---- service/src/secure-startup.test.ts | 11 ++ service/src/secure-startup.ts | 6 ++ 8 files changed, 397 insertions(+), 77 deletions(-) diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index eb9ec7d0..dd3f54b5 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -1,4 +1,5 @@ export const BRIDGE_PROTOCOL_VERSION = 1 as const; +export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; @@ -95,3 +96,7 @@ export class BridgeProtocolError extends Error { export function bridgeWorkerPath(workerId: string): string { return `/bridge/workers/${encodeURIComponent(workerId)}`; } + +export function isValidBridgeWorkerId(workerId: string): boolean { + return BRIDGE_WORKER_ID_PATTERN.test(workerId); +} diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 17451aa8..f4b73c88 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -753,6 +753,7 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', test('worker preserves a definite rejection when its heartbeat fails', async () => { let registrations = 0; let rejectedSettlement = false; + let settlementAttempts = 0; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', @@ -768,7 +769,9 @@ test('worker preserves a definite rejection when its heartbeat fails', async () fetchImpl: async (input, init) => { if (String(input).endsWith('/workers/register')) { registrations += 1; - if (registrations > 1) throw new TypeError('registration unavailable'); + if (registrations === 2) { + throw new TypeError('registration unavailable'); + } return new Response( JSON.stringify({ protocolVersion: 1, @@ -787,8 +790,15 @@ test('worker preserves a definite rejection when its heartbeat fails', async () headers: { 'Content-Type': 'application/json' }, }); } + settlementAttempts += 1; rejectedSettlement = JSON.parse(String(init?.body) || '{}').status === 'rejected'; + if (settlementAttempts === 1) { + return new Response(JSON.stringify({ error: 'unavailable' }), { + status: 503, + headers: { 'Content-Type': 'application/json' }, + }); + } return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), { status: 200, headers: { 'Content-Type': 'application/json' } }, @@ -810,8 +820,9 @@ test('worker preserves a definite rejection when its heartbeat fails', async () request: { body: { language: 'bash' }, headers: {} }, }); - assert.ok(registrations > 1); + assert.ok(registrations >= 3); assert.equal(rejectedSettlement, true); + assert.equal(settlementAttempts, 2); }); test('worker quarantines a stateful workspace after a sandbox 5xx response', async () => { @@ -1234,6 +1245,73 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( } }); +test('worker rejects an assignment after ambiguous acknowledgement delivery', async () => { + let rejectedSettlement = false; + let acknowledgementAttempts = 0; + 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: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + rejectionAckGraceMs: 500, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/ack')) { + acknowledgementAttempts += 1; + throw new TypeError('acknowledgement response lost'); + } + if (String(input).endsWith('/settle')) { + rejectedSettlement = + JSON.parse(String(init?.body) || '{}').status === 'rejected'; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + if (String(input).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' } }, + ); + } + return new Response( + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 0, + assignment: { + protocolVersion: 1, + assignmentId: 'ambiguous-ack', + 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(), + remainingMs: 1_000, + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await assert.rejects(worker.lease(), /acknowledgement response lost/); + assert.equal(acknowledgementAttempts, 1); + assert.equal(rejectedSettlement, true); +}); + test('worker clamps rejected settlement errors to the protocol limit', async () => { let rejection = ''; const worker = new BridgeWorker({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index b86afa68..cfac6a20 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -45,6 +45,7 @@ const DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS = 10_000; const DEFAULT_CANCELLATION_POLL_INTERVAL_MS = 500; const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; +const REGISTRATION_RETRY_DELAY_MS = 100; const SETTLEMENT_RETRY_DELAY_MS = 100; const REJECTION_ACK_GRACE_MS = 30_000; const MAX_SETTLEMENT_ERROR_LENGTH = 4_096; @@ -234,52 +235,47 @@ export class BridgeWorker { ), }; const acknowledgementStartedAtMs = Date.now(); - await this.timedRequest( - this.assignmentUrl(adjustedAssignment, 'ack'), - { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - incarnationId: this.incarnationId, - generation: adjustedAssignment.generation, - leaseToken: adjustedAssignment.leaseToken, - }, - Math.max( - 1, - this.options.leaseAckTransportTimeoutMs ?? - DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, - ), - signal, - ); + try { + await this.timedRequest( + this.assignmentUrl(adjustedAssignment, 'ack'), + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + generation: adjustedAssignment.generation, + leaseToken: adjustedAssignment.leaseToken, + }, + Math.max( + 1, + this.options.leaseAckTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + ), + signal, + ); + } catch (error) { + const definiteRejection = + error instanceof BridgeProtocolError && + error.status != null && + error.status < 500 && + error.status !== 408 && + error.status !== 429; + if (!definiteRejection) { + await this.rejectUnexecutedAssignment( + adjustedAssignment, + 'Bridge lease acknowledgement delivery was ambiguous', + ); + } + throw error; + } const remainingMs = Math.max( 0, (adjustedAssignment.remainingMs ?? 0) - (Date.now() - acknowledgementStartedAtMs), ); if (remainingMs <= 0) { - const heartbeatController = new AbortController(); - const heartbeat = this.maintainRegistration( - heartbeatController.signal, - ).catch(() => undefined); - try { - await this.settleWithRetry( - adjustedAssignment, - { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - generation: adjustedAssignment.generation, - leaseToken: adjustedAssignment.leaseToken, - incarnationId: this.incarnationId, - status: 'rejected', - error: 'Bridge assignment expired during lease acknowledgement', - }, - Date.now() + - Math.max( - 0, - this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, - ), - ); - } finally { - heartbeatController.abort(); - await heartbeat; - } + await this.rejectUnexecutedAssignment( + adjustedAssignment, + 'Bridge assignment expired during lease acknowledgement', + ); throw new BridgeProtocolError( 'Bridge assignment expired during lease acknowledgement', ); @@ -435,18 +431,36 @@ export class BridgeWorker { assignment.runtimeSessionId != null && settlement.status === 'rejected' && sandboxRejectedExecution; - await this.settleWithRetry( - assignment, - settlement, - localDeadlineAtMs + - (knownCleanStatefulRejection - ? Math.max( + if (knownCleanStatefulRejection) { + heartbeatController.abort(); + await heartbeat; + const recoveryHeartbeatController = new AbortController(); + const recoveryHeartbeat = this.maintainRegistration( + recoveryHeartbeatController.signal, + true, + ).catch(() => undefined); + try { + await this.settleWithRetry( + assignment, + settlement, + localDeadlineAtMs + + Math.max( 0, this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, - ) - : 0), - knownCleanStatefulRejection ? undefined : signal, - ); + ), + ); + } finally { + recoveryHeartbeatController.abort(); + await recoveryHeartbeat; + } + } else { + await this.settleWithRetry( + assignment, + settlement, + localDeadlineAtMs, + signal, + ); + } } finally { heartbeatController.abort(); await heartbeat; @@ -502,6 +516,7 @@ export class BridgeWorker { private async maintainRegistration( signal: AbortSignal, + retryTransient = false, ): Promise { while (!signal.aborted) { const heartbeatIntervalMs = Math.max( @@ -516,7 +531,50 @@ export class BridgeWorker { signal, ); if (signal.aborted) return; - await this.register(signal); + try { + await this.register(signal); + } catch (error) { + const terminal = + error instanceof BridgeProtocolError && + (error.status === 401 || + error.status === 403 || + error.code === 'WORKER_FENCED' || + error.code === 'WORKER_QUARANTINED'); + if (!retryTransient || terminal || signal.aborted) throw error; + await this.delay(REGISTRATION_RETRY_DELAY_MS, signal); + } + } + } + + private async rejectUnexecutedAssignment( + assignment: BridgeAssignment, + error: string, + ): Promise { + const heartbeatController = new AbortController(); + const heartbeat = this.maintainRegistration( + heartbeatController.signal, + true, + ).catch(() => undefined); + try { + await this.settleWithRetry( + assignment, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, + status: 'rejected', + error, + }, + Date.now() + + Math.max( + 0, + this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, + ), + ); + } finally { + heartbeatController.abort(); + await heartbeat; } } diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 2d25a1b0..00596007 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -5,12 +5,14 @@ import type { NextFunction, Request, RequestHandler, Response } from 'express'; import type { BridgeWorkerRegistration } from '../../../packages/code/src/protocol'; import type { CodeBridgeAssignment, CodeBridgeSettlement } from './store'; -import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { + BRIDGE_PROTOCOL_VERSION, + isValidBridgeWorkerId, +} 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; @@ -43,7 +45,7 @@ function bridgeAuth(req: Request, res: Response, next: NextFunction): void { } function validWorkerId(value: string): boolean { - return WORKER_ID_PATTERN.test(value); + return isValidBridgeWorkerId(value); } function validIncarnationId(value: unknown): value is string { diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index afa44526..c8a224c8 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -161,6 +161,15 @@ describe('RedisBridgeStore', () => { await expect(completion).resolves.toMatchObject({ status: 'rejected' }); }); + test('bounds a stalled Redis lease claim', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.eval = (() => new Promise(() => undefined)) as Redis['eval']; + + await expect( + timedStore.lease('stalled-worker', incarnationId, 0), + ).rejects.toThrow('Bridge lease claim timed out'); + }); + test('bounds a stalled quarantine command', async () => { const timedStore = new RedisBridgeStore(redis, 60, 10); redis.eval = (() => new Promise(() => undefined)) as Redis['eval']; @@ -299,6 +308,59 @@ describe('RedisBridgeStore', () => { }); }); + test('preserves a workspace fence when an acknowledged lease expires', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'ack-expired-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'ack-expired-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-ack-expired', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'ack-expired-worker', + incarnationId, + 1_000, + ); + await store.acknowledgeLease( + 'ack-expired-worker', + incarnationId, + assignment?.assignmentId ?? '', + assignment?.generation ?? 0, + assignment?.leaseToken ?? '', + ); + const storedKey = `codeapi:bridge:v1:assignment:${assignment?.assignmentId}`; + const stored = JSON.parse( + (await redis.get(storedKey)) ?? '{}', + ) as Record; + stored.expiresAt = new Date(0).toISOString(); + await redis.set(storedKey, JSON.stringify(stored), 'EX', 30); + + await expect( + store.lease('ack-expired-worker', incarnationId, 0), + ).resolves.toBeUndefined(); + expect( + await redis.keys( + 'codeapi:bridge:v1:worker:ack-expired-worker:workspace:*:quarantined', + ), + ).toHaveLength(1); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('clears a workspace fence when dispatch cancels before lease', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -383,6 +445,46 @@ describe('RedisBridgeStore', () => { }); }); + test('restores queue expiry when returning a lease', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'returned-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'returned-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'returned-worker', + incarnationId, + 1_000, + ); + await store.returnLease(assignment!); + + expect( + await redis.ttl( + `codeapi:bridge:v1:worker:returned-worker:incarnation:${incarnationId}:assignments`, + ), + ).toBeGreaterThan(0); + expect( + await store.lease('returned-worker', incarnationId, 1_000), + ).toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('returns a popped assignment after a transient Redis read failure', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index c30f74c8..2742a56e 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -203,6 +203,19 @@ export class RedisBridgeStore { } } + private async leaseCommand( + command: Promise, + signal: AbortSignal | undefined, + label: string, + ): Promise { + return await boundedCommand( + command, + this.redisCommandTimeoutMs, + label, + signal, + ); + } + async register(registration: BridgeWorkerRegistration): Promise { const script = [ 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', @@ -446,7 +459,17 @@ export class RedisBridgeStore { (firstPoll || Date.now() < deadline) ) { firstPoll = false; - const assignmentId = await this.claimOrPopLease(workerId, incarnationId); + let assignmentId: string | null; + try { + assignmentId = await this.leaseCommand( + this.claimOrPopLease(workerId, incarnationId), + signal, + 'Bridge lease claim', + ); + } catch (error) { + if (signalAborted(signal)) return undefined; + throw error; + } if (assignmentId == null) { await delay( Math.min(POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), @@ -455,20 +478,32 @@ export class RedisBridgeStore { continue; } try { - const assignment = await this.readAssignment(assignmentId); + const assignment = await this.leaseCommand( + this.readAssignment(assignmentId), + signal, + 'Bridge lease assignment read', + ); if ( assignment == null || assignment.workerId !== workerId || assignment.incarnationId !== incarnationId ) { - await this.discardLeaseClaim(workerId, incarnationId, assignmentId); + await this.leaseCommand( + this.discardLeaseClaim(workerId, incarnationId, assignmentId), + signal, + 'Bridge lease claim discard', + ); continue; } if (signalAborted(signal)) { await this.returnLease(assignment); return undefined; } - const registration = await this.registration(workerId); + const registration = await this.leaseCommand( + this.registration(workerId), + signal, + 'Bridge lease registration read', + ); if (registration?.incarnationId !== incarnationId) { throw new BridgeStoreError( 'WORKER_FENCED', @@ -476,8 +511,24 @@ export class RedisBridgeStore { ); } if (Date.parse(assignment.expiresAt) <= Date.now()) { - await this.clearUndeliveredWorkspaceFence(assignment); - await this.discardLeaseClaim(workerId, incarnationId, assignmentId); + const acknowledged = + (await this.leaseCommand( + this.redis.get(leaseAckKey(workerId, incarnationId)), + signal, + 'Bridge lease acknowledgement read', + )) === assignmentId; + if (!acknowledged) { + await this.leaseCommand( + this.clearUndeliveredWorkspaceFence(assignment), + signal, + 'Bridge undelivered workspace recovery', + ); + } + await this.leaseCommand( + this.discardLeaseClaim(workerId, incarnationId, assignmentId), + signal, + 'Bridge expired lease discard', + ); continue; } if (signalAborted(signal)) { @@ -498,6 +549,7 @@ export class RedisBridgeStore { incarnationId, assignmentId, ); + if (signalAborted(signal)) return undefined; throw error; } } @@ -602,21 +654,27 @@ export class RedisBridgeStore { incarnationId: string, assignmentId: string, ): Promise { - await this.redis.eval( - [ - "if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end", - "if redis.call('GET', KEYS[3]) ~= ARGV[1] then return 0 end", - "redis.call('DEL', KEYS[3], KEYS[4])", - "redis.call('LREM', KEYS[2], 0, ARGV[1])", - "redis.call('LPUSH', KEYS[2], ARGV[1])", - 'return 1', - ].join('\n'), - 4, - assignmentKey(assignmentId), - queueKey(workerId, incarnationId), - leaseClaimKey(workerId, incarnationId), - leaseAckKey(workerId, incarnationId), - assignmentId, + await boundedCommand( + this.redis.eval( + [ + "if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end", + "if redis.call('GET', KEYS[3]) ~= ARGV[1] then return 0 end", + "local ttl = redis.call('TTL', KEYS[1])", + "redis.call('DEL', KEYS[3], KEYS[4])", + "redis.call('LREM', KEYS[2], 0, ARGV[1])", + "redis.call('LPUSH', KEYS[2], ARGV[1])", + "if ttl > 0 then redis.call('EXPIRE', KEYS[2], ttl) end", + 'return 1', + ].join('\n'), + 4, + assignmentKey(assignmentId), + queueKey(workerId, incarnationId), + leaseClaimKey(workerId, incarnationId), + leaseAckKey(workerId, incarnationId), + assignmentId, + ), + this.redisCommandTimeoutMs, + 'Bridge lease return', ); } diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index c2e3c1f5..0ddfeae3 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -342,6 +342,17 @@ describe('sandbox backend policy', () => { expect(() => validateApiBridgePolicy()).not.toThrow(); }); + test('API bridge policy rejects worker IDs the router cannot accept', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.BRIDGE_WORKER_ID = 'engineering/vm'; + env.BRIDGE_TOKEN = 'development-bridge-token'; + env.PTC_MODE = 'replay'; + + expect(() => validateApiBridgePolicy()).toThrow( + 'must match the bridge worker ID format', + ); + }); + test('remote bridge requires a positive finite job timeout', () => { env.SANDBOX_BACKEND = 'remote-bridge'; env.BRIDGE_WORKER_ID = 'engineering-vm'; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 437a6276..357a2977 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -4,6 +4,7 @@ import { lambdaMicrovmNumericConfigError, } from './config'; import { INTERNAL_SERVICE_TOKEN_ENV } from './internal-service-auth'; +import { isValidBridgeWorkerId } from '../../packages/code/src/protocol'; export class SecureStartupConfigError extends Error { constructor(message: string) { @@ -58,6 +59,11 @@ export function validateApiBridgePolicy(): void { if (env.SANDBOX_BACKEND !== 'remote-bridge') return; requireSafeWholeNumber('JOB_TIMEOUT', env.JOB_TIMEOUT, 1); requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); + if (!isValidBridgeWorkerId(env.BRIDGE_WORKER_ID ?? '')) { + throw new SecureStartupConfigError( + 'CODEAPI_BRIDGE_WORKER_ID must match the bridge worker ID format', + ); + } if (env.HARDENED_SANDBOX_MODE) { requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); } else { From 4bfedb91e04678204e648f48e049292cc746ea01 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 13:49:55 -0400 Subject: [PATCH 22/29] fix: bound bridge control lifetimes --- service/src/bridge/store.test.ts | 54 ++++++++++++++++++++++++++++++++ service/src/bridge/store.ts | 48 ++++++++++++++++++---------- 2 files changed, 86 insertions(+), 16 deletions(-) diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index c8a224c8..330f2cb1 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -170,6 +170,60 @@ describe('RedisBridgeStore', () => { ).rejects.toThrow('Bridge lease claim timed out'); }); + test('bounds a stalled Redis worker registration', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.eval = (() => new Promise(() => undefined)) as Redis['eval']; + + await expect( + timedStore.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'stalled-registration-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toThrow('Bridge worker registration timed out'); + }); + + test('retains cancellation through the assignment lifetime', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'cancel-ttl-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'cancel-ttl-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 120_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'cancel-ttl-worker', + incarnationId, + 1_000, + ); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + + expect( + await redis.ttl( + `codeapi:bridge:v1:assignment:${assignment?.assignmentId}:cancelled`, + ), + ).toBeGreaterThan(30); + }); + test('bounds a stalled quarantine command', async () => { const timedStore = new RedisBridgeStore(redis, 60, 10); redis.eval = (() => new Promise(() => undefined)) as Redis['eval']; diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 2742a56e..ff1ee4b7 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -236,19 +236,23 @@ export class RedisBridgeStore { '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:`, + await boundedCommand( + 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:`, + ), + this.redisCommandTimeoutMs, + 'Bridge worker registration', ), ); if (result === -2) { @@ -990,9 +994,21 @@ export class RedisBridgeStore { ); } - private async cancel(assignmentId: string): Promise { + private async cancel( + assignmentId: string, + assignment?: StoredAssignment, + ): Promise { + const ttlSeconds = + assignment == null + ? 30 + : assignmentTtlSeconds(Date.parse(assignment.expiresAt)); await boundedCommand( - this.redis.set(cancellationKey(assignmentId), '1', 'EX', 30), + this.redis.set( + cancellationKey(assignmentId), + '1', + 'EX', + ttlSeconds, + ), this.redisCommandTimeoutMs, 'Bridge assignment cancellation', ); @@ -1075,7 +1091,7 @@ export class RedisBridgeStore { assignment: StoredAssignment | undefined, ): Promise { await Promise.all([ - this.cancel(assignmentId), + this.cancel(assignmentId, assignment), assignment == null ? boundedCommand( this.releaseLock(workerId, assignmentId), From 891882cd76ae4458f1275447886ed1ea2dd075e7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 14:03:14 -0400 Subject: [PATCH 23/29] fix: retain bridge recovery ownership --- service/src/bridge/router.ts | 24 ++++++++++++++++------ service/src/bridge/store.test.ts | 32 +++++++++++++++++++++++++++++ service/src/bridge/store.ts | 35 ++++++++++++++++++++++++-------- 3 files changed, 76 insertions(+), 15 deletions(-) diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 00596007..12b64762 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -334,12 +334,24 @@ router.post( 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 }); + const cancellationController = new AbortController(); + const abortCancellation = (): void => cancellationController.abort(); + req.once('aborted', abortCancellation); + res.once('close', abortCancellation); + try { + const cancelled = await bridgeStore.cancelled( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + cancellationController.signal, + ); + if (!cancellationController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); + } + } finally { + req.off('aborted', abortCancellation); + res.off('close', abortCancellation); + } }), ); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 330f2cb1..635f9534 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -188,6 +188,19 @@ describe('RedisBridgeStore', () => { ).rejects.toThrow('Bridge worker registration timed out'); }); + test('bounds stalled Redis reads during cancellation polling', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.get = (() => new Promise(() => undefined)) as Redis['get']; + + await expect( + timedStore.cancelled( + 'stalled-cancellation-worker', + incarnationId, + 'assignment-stalled-cancellation', + ), + ).rejects.toThrow('Bridge cancellation assignment read timed out'); + }); + test('retains cancellation through the assignment lifetime', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -1294,6 +1307,25 @@ describe('RedisBridgeStore', () => { await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED', }); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'lost-worker', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + expect( + await redis.get('codeapi:bridge:v1:worker:lost-worker:lock'), + ).toBe(assignment?.assignmentId ?? null); + await redis.del( + 'codeapi:bridge:v1:worker:lost-worker:lock', + 'codeapi:bridge:v1:worker:lost-worker:lock:incarnation', + ); await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'lost-worker', diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index ff1ee4b7..de9e4833 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -830,9 +830,18 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, assignmentId: string, + signal?: AbortSignal, ): Promise { - const assignment = await this.readAssignment(assignmentId); - const registration = await this.registration(workerId); + const assignment = await this.leaseCommand( + this.readAssignment(assignmentId), + signal, + 'Bridge cancellation assignment read', + ); + const registration = await this.leaseCommand( + this.registration(workerId), + signal, + 'Bridge cancellation registration read', + ); if ( assignment == null || assignment.workerId !== workerId || @@ -841,7 +850,13 @@ export class RedisBridgeStore { ) { return true; } - return (await this.redis.exists(cancellationKey(assignmentId))) === 1; + return ( + (await this.leaseCommand( + this.redis.exists(cancellationKey(assignmentId)), + signal, + 'Bridge cancellation marker read', + )) === 1 + ); } async quarantine( @@ -1190,12 +1205,12 @@ export class RedisBridgeStore { " redis.call('DEL', KEYS[3], KEYS[4])", 'end', 'if queued == 0 and acknowledged and ARGV[2] == "1" and redis.call(\'GET\', KEYS[5]) == ARGV[1] then', - ' return 0', + ' return -1', 'end', "return redis.call('DEL', KEYS[1], KEYS[3], KEYS[4])", ].join('\n'); - await Promise.all([ - boundedCommand( + const cleanupResult = Number( + await boundedCommand( this.redis.eval( cleanupScript, keys.length, @@ -1206,12 +1221,14 @@ export class RedisBridgeStore { this.redisCommandTimeoutMs, 'Bridge assignment cleanup', ), - boundedCommand( + ); + if (cleanupResult !== -1) { + await boundedCommand( this.releaseLock(assignment.workerId, assignment.assignmentId), this.redisCommandTimeoutMs, 'Bridge assignment lock release', - ), - ]); + ); + } } private async releaseLock( From c4c4cc6026fcb5284559d6f04546504939eac4d9 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 14:24:05 -0400 Subject: [PATCH 24/29] fix: isolate bridge control state --- packages/code/src/worker.test.ts | 45 ++++++++++++++ packages/code/src/worker.ts | 7 ++- service/src/bridge/router.ts | 59 ++++++++++++------ service/src/bridge/store.test.ts | 74 ++++++++++++++++++++++ service/src/bridge/store.ts | 102 ++++++++++++++++++++----------- 5 files changed, 232 insertions(+), 55 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index f4b73c88..94945cfd 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -872,6 +872,51 @@ test('worker quarantines a stateful workspace after a sandbox 5xx response', asy assert.equal(settlementAttempted, false); }); +test('worker treats a non-JSON sandbox 4xx as a definite rejection', async () => { + let rejectedSettlement = false; + 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: async (input, init) => { + if (String(input).endsWith('/execute')) { + return new Response('not found', { + status: 404, + headers: { 'Content-Type': 'text/html' }, + }); + } + rejectedSettlement = + JSON.parse(String(init?.body) || '{}').status === 'rejected'; + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'non-json-404', + 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(), + remainingMs: 1_000, + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }); + assert.equal(rejectedSettlement, true); +}); + test('worker quarantines a stateful workspace after the sandbox request aborts', async () => { let settlementAttempted = false; const fetchImpl: typeof fetch = async (input, init) => { diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index cfac6a20..4fc4a78c 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -373,7 +373,12 @@ export class BridgeWorker { signal: executionController.signal, }, ); - const payload = (await response.json()) as object; + let payload: object = {}; + try { + payload = (await response.json()) as object; + } catch (error) { + if (response.ok) throw error; + } if (!response.ok) { sandboxRejectedExecution = response.status >= 400 && diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 12b64762..910a8dd4 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -277,14 +277,27 @@ router.post( return; } try { - await bridgeStore.acknowledgeLease( - req.params.workerId, - body.incarnationId, - req.params.assignmentId, - Number(body.generation), - body.leaseToken, - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, accepted: true }); + const acknowledgementController = new AbortController(); + const abortAcknowledgement = (): void => + acknowledgementController.abort(); + req.once('aborted', abortAcknowledgement); + res.once('close', abortAcknowledgement); + try { + await bridgeStore.acknowledgeLease( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + Number(body.generation), + body.leaseToken, + acknowledgementController.signal, + ); + if (!acknowledgementController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, accepted: true }); + } + } finally { + req.off('aborted', abortAcknowledgement); + res.off('close', abortAcknowledgement); + } } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); @@ -304,15 +317,27 @@ router.post( return; } try { - await bridgeStore.settle( - req.params.workerId, - req.params.assignmentId, - settlement, - ); - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - accepted: true, - }); + const settlementController = new AbortController(); + const abortSettlement = (): void => settlementController.abort(); + req.once('aborted', abortSettlement); + res.once('close', abortSettlement); + try { + await bridgeStore.settle( + req.params.workerId, + req.params.assignmentId, + settlement, + settlementController.signal, + ); + if (!settlementController.signal.aborted) { + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + accepted: true, + }); + } + } finally { + req.off('aborted', abortSettlement); + res.off('close', abortSettlement); + } } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 635f9534..da2561c6 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -201,6 +201,80 @@ describe('RedisBridgeStore', () => { ).rejects.toThrow('Bridge cancellation assignment read timed out'); }); + test('bounds stalled Redis reads during lease acknowledgement', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.get = (() => new Promise(() => undefined)) as Redis['get']; + + await expect( + timedStore.acknowledgeLease( + 'stalled-ack-worker', + incarnationId, + 'assignment-stalled-ack', + 1, + 'lease-token-that-is-long-enough-for-testing', + ), + ).rejects.toThrow('Bridge acknowledgement assignment read timed out'); + }); + + test('bounds stalled Redis reads during settlement', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.get = (() => new Promise(() => undefined)) as Redis['get']; + + await expect( + timedStore.settle('stalled-settlement-worker', 'assignment-stalled', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + incarnationId, + status: 'rejected', + error: 'test', + }), + ).rejects.toThrow('Bridge settlement existing read timed out'); + }); + + test('encodes worker IDs so Redis key families cannot collide', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'foo', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'foo', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('foo', incarnationId, 1_000); + + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'foo:lock', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + expect( + await redis.get('codeapi:bridge:v1:worker:foo:lock'), + ).toBe(assignment?.assignmentId ?? null); + expect( + await redis.get('codeapi:bridge:v1:worker:foo%3Alock'), + ).not.toBeNull(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('retains cancellation through the assignment lifetime', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index de9e4833..e8f8aaa8 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -48,19 +48,19 @@ interface StoredAssignment extends CodeBridgeAssignment { } function workerKey(workerId: string): string { - return `${PREFIX}:worker:${workerId}`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}`; } function workerIncarnationKey(workerId: string): string { - return `${PREFIX}:worker:${workerId}:incarnation`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation`; } function incarnationFenceKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:fenced`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:fenced`; } function quarantineKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:quarantined`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:quarantined`; } function workspaceQuarantineKey( @@ -70,31 +70,31 @@ function workspaceQuarantineKey( const sessionHash = createHash('sha256') .update(runtimeSessionId) .digest('hex'); - return `${PREFIX}:worker:${workerId}:workspace:${sessionHash}:quarantined`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:workspace:${sessionHash}:quarantined`; } function queueKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:assignments`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:assignments`; } function leaseClaimKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:lease-claim`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-claim`; } function leaseAckKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${workerId}:incarnation:${incarnationId}:lease-ack`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-ack`; } function generationKey(workerId: string): string { - return `${PREFIX}:worker:${workerId}:generation`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:generation`; } function lockKey(workerId: string): string { - return `${PREFIX}:worker:${workerId}:lock`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:lock`; } function lockIncarnationKey(workerId: string): string { - return `${PREFIX}:worker:${workerId}:lock:incarnation`; + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:lock:incarnation`; } function assignmentKey(assignmentId: string): string { @@ -249,7 +249,7 @@ export class RedisBridgeStore { registration.incarnationId, JSON.stringify(registration), String(this.workerTtlSeconds), - `${PREFIX}:worker:${registration.workerId}:incarnation:`, + `${PREFIX}:worker:${encodeURIComponent(registration.workerId)}:incarnation:`, ), this.redisCommandTimeoutMs, 'Bridge worker registration', @@ -566,9 +566,18 @@ export class RedisBridgeStore { assignmentId: string, generation: number, leaseToken: string, + signal?: AbortSignal, ): Promise { - const assignment = await this.readAssignment(assignmentId); - const registration = await this.registration(workerId); + const assignment = await this.leaseCommand( + this.readAssignment(assignmentId), + signal, + 'Bridge acknowledgement assignment read', + ); + const registration = await this.leaseCommand( + this.registration(workerId), + signal, + 'Bridge acknowledgement registration read', + ); if ( assignment == null || assignment.workerId !== workerId || @@ -584,17 +593,21 @@ export class RedisBridgeStore { } const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); const acknowledged = Number( - await this.redis.eval( - [ - "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", - "redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2])", - 'return 1', - ].join('\n'), - 2, - leaseClaimKey(workerId, incarnationId), - leaseAckKey(workerId, incarnationId), - assignmentId, - String(ttlSeconds), + await this.leaseCommand( + this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", + "redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2])", + 'return 1', + ].join('\n'), + 2, + leaseClaimKey(workerId, incarnationId), + leaseAckKey(workerId, incarnationId), + assignmentId, + String(ttlSeconds), + ), + signal, + 'Bridge lease acknowledgement', ), ); if (acknowledged !== 1) { @@ -724,10 +737,13 @@ export class RedisBridgeStore { workerId: string, assignmentId: string, settlement: CodeBridgeSettlement, + signal?: AbortSignal, ): Promise { const serializedSettlement = JSON.stringify(settlement); - const existingSettlement = await this.redis.get( - settlementKey(assignmentId), + const existingSettlement = await this.leaseCommand( + this.redis.get(settlementKey(assignmentId)), + signal, + 'Bridge settlement existing read', ); if (existingSettlement === serializedSettlement) return; if (existingSettlement != null) { @@ -736,7 +752,11 @@ export class RedisBridgeStore { 'Bridge assignment was already settled with a different result', ); } - const assignment = await this.readAssignment(assignmentId); + const assignment = await this.leaseCommand( + this.readAssignment(assignmentId), + signal, + 'Bridge settlement assignment read', + ); if (assignment == null) { throw new BridgeStoreError( 'ASSIGNMENT_NOT_FOUND', @@ -749,7 +769,11 @@ export class RedisBridgeStore { 'Bridge assignment belongs to another worker', ); } - const registration = await this.registration(workerId); + const registration = await this.leaseCommand( + this.registration(workerId), + signal, + 'Bridge settlement registration read', + ); if ( settlement.incarnationId !== assignment.incarnationId || registration?.incarnationId !== settlement.incarnationId || @@ -796,14 +820,18 @@ export class RedisBridgeStore { 'return 1', ].join('\n'); const accepted = Number( - await this.redis.eval( - script, - settlementKeys.length, - ...settlementKeys, - serializedSettlement, - String(ttlSeconds), - assignmentId, - settlement.status, + await this.leaseCommand( + this.redis.eval( + script, + settlementKeys.length, + ...settlementKeys, + serializedSettlement, + String(ttlSeconds), + assignmentId, + settlement.status, + ), + signal, + 'Bridge settlement commit', ), ); if (accepted === -1) { From 96177c30d99e48da3b6cf1f7f25a4c2f5c734e9d Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 14:44:09 -0400 Subject: [PATCH 25/29] fix: bound bridge workspace reset --- service/src/bridge/router.ts | 24 ++++++++++++++++++------ service/src/bridge/store.test.ts | 13 +++++++++++++ service/src/bridge/store.ts | 29 +++++++++++++++++------------ 3 files changed, 48 insertions(+), 18 deletions(-) diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 910a8dd4..3f03a533 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -186,12 +186,24 @@ router.post( return; } try { - await bridgeStore.resetWorkspace( - workerId, - body.incarnationId, - body.runtimeSessionId, - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); + const resetController = new AbortController(); + const abortReset = (): void => resetController.abort(); + req.once('aborted', abortReset); + res.once('close', abortReset); + try { + await bridgeStore.resetWorkspace( + workerId, + body.incarnationId, + body.runtimeSessionId, + resetController.signal, + ); + if (!resetController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); + } + } finally { + req.off('aborted', abortReset); + res.off('close', abortReset); + } } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index da2561c6..e28b32e3 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -232,6 +232,19 @@ describe('RedisBridgeStore', () => { ).rejects.toThrow('Bridge settlement existing read timed out'); }); + test('bounds a stalled Redis workspace reset', async () => { + const timedStore = new RedisBridgeStore(redis, 60, 10); + redis.eval = (() => new Promise(() => undefined)) as Redis['eval']; + + await expect( + timedStore.resetWorkspace( + 'stalled-reset-worker', + incarnationId, + 'rt-stalled-reset', + ), + ).rejects.toThrow('Bridge workspace reset timed out'); + }); + test('encodes worker IDs so Redis key families cannot collide', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index e8f8aaa8..0b2a84e5 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -925,20 +925,25 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, runtimeSessionId: string, + signal?: AbortSignal, ): Promise { const result = Number( - await this.redis.eval( - [ - "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end", - "if redis.call('EXISTS', KEYS[2]) == 1 then return -2 end", - "redis.call('DEL', KEYS[3])", - 'return 1', - ].join('\n'), - 3, - workerIncarnationKey(workerId), - lockKey(workerId), - workspaceQuarantineKey(workerId, runtimeSessionId), - incarnationId, + await this.leaseCommand( + this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end", + "if redis.call('EXISTS', KEYS[2]) == 1 then return -2 end", + "redis.call('DEL', KEYS[3])", + 'return 1', + ].join('\n'), + 3, + workerIncarnationKey(workerId), + lockKey(workerId), + workspaceQuarantineKey(workerId, runtimeSessionId), + incarnationId, + ), + signal, + 'Bridge workspace reset', ), ); if (result === -1) { From 1530105aeee886ca0ab5f148f4fc9dc5f5c809f4 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 14:59:30 -0400 Subject: [PATCH 26/29] fix: validate bridge deployment inputs --- packages/code/src/cli.test.ts | 26 ++++++++++++++++++++++++++ packages/code/src/cli.ts | 9 ++++++++- packages/code/src/protocol.test.ts | 9 ++++++++- service/Dockerfile.node | 2 ++ service/Dockerfile.service | 1 + service/src/secure-startup.test.ts | 11 +++++++++++ service/src/secure-startup.ts | 5 +++++ 7 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 packages/code/src/cli.test.ts diff --git a/packages/code/src/cli.test.ts b/packages/code/src/cli.test.ts new file mode 100644 index 00000000..837e9e99 --- /dev/null +++ b/packages/code/src/cli.test.ts @@ -0,0 +1,26 @@ +import assert from 'node:assert/strict'; +import { spawnSync } from 'node:child_process'; +import { fileURLToPath } from 'node:url'; +import test from 'node:test'; + +test('CLI rejects an invalid worker ID before entering the run loop', () => { + const result = spawnSync( + process.execPath, + [fileURLToPath(new URL('./cli.js', import.meta.url))], + { + encoding: 'utf8', + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'https://code.example/v1', + LIBRECHAT_CODE_WORKER_TOKEN: 'worker-secret', + LIBRECHAT_CODE_WORKER_ID: 'engineering/vm', + }, + }, + ); + + assert.notEqual(result.status, 0); + assert.match( + result.stderr, + /LIBRECHAT_CODE_WORKER_ID must match the bridge worker ID format/, + ); +}); diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index a10e2ed4..a7f60cd7 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -1,6 +1,7 @@ #!/usr/bin/env node import { createHash } from 'node:crypto'; import { BridgeWorker } from './worker.js'; +import { isValidBridgeWorkerId } from './protocol.js'; function required(name: string): string { const value = process.env[name]?.trim(); @@ -32,10 +33,16 @@ if (statefulWorkspace && !sandboxEndpoint.includes('{runtimeSessionId}')) { 'LIBRECHAT_CODE_STATEFUL_WORKSPACE requires LIBRECHAT_CODE_SANDBOX_ENDPOINT to contain {runtimeSessionId}', ); } +const workerId = required('LIBRECHAT_CODE_WORKER_ID'); +if (!isValidBridgeWorkerId(workerId)) { + throw new Error( + 'LIBRECHAT_CODE_WORKER_ID must match the bridge worker ID format', + ); +} const worker = new BridgeWorker({ codeApiUrl: required('LIBRECHAT_CODE_URL'), token: required('LIBRECHAT_CODE_WORKER_TOKEN'), - workerId: required('LIBRECHAT_CODE_WORKER_ID'), + workerId, sandboxEndpoint, capabilities: { statefulWorkspace, diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts index 9223b3ab..38f859e2 100644 --- a/packages/code/src/protocol.test.ts +++ b/packages/code/src/protocol.test.ts @@ -1,6 +1,6 @@ import assert from 'node:assert/strict'; import test from 'node:test'; -import { bridgeWorkerPath } from './protocol.js'; +import { bridgeWorkerPath, isValidBridgeWorkerId } from './protocol.js'; test('bridgeWorkerPath encodes worker-controlled path segments', () => { assert.equal( @@ -8,3 +8,10 @@ test('bridgeWorkerPath encodes worker-controlled path segments', () => { '/bridge/workers/vm%2Fexample%20worker', ); }); + +test('bridge worker IDs reject path, whitespace, and oversized values', () => { + assert.equal(isValidBridgeWorkerId('engineering-vm:1'), true); + assert.equal(isValidBridgeWorkerId('engineering/vm'), false); + assert.equal(isValidBridgeWorkerId('engineering vm'), false); + assert.equal(isValidBridgeWorkerId('a'.repeat(129)), false); +}); diff --git a/service/Dockerfile.node b/service/Dockerfile.node index 46f4c319..8a6006e4 100644 --- a/service/Dockerfile.node +++ b/service/Dockerfile.node @@ -8,6 +8,7 @@ RUN npm ci FROM base AS builder COPY service/src ./src COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ RUN npx tsc -p tsconfig.json @@ -26,6 +27,7 @@ FROM base AS development ENV NODE_ENV=development COPY service/src ./src COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ RUN npm install -g ts-node typescript EXPOSE 3000 9230 diff --git a/service/Dockerfile.service b/service/Dockerfile.service index 53931e71..3a89dc15 100644 --- a/service/Dockerfile.service +++ b/service/Dockerfile.service @@ -17,6 +17,7 @@ FROM base AS builder COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY shared /usr/src/shared +COPY packages/code/src /usr/src/packages/code/src COPY service/tsconfig.json ./ RUN bun build ./src/service-api.ts --outdir .build --target bun --external '@opentelemetry/*' diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 0ddfeae3..a365f355 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -353,6 +353,17 @@ describe('sandbox backend policy', () => { ); }); + test('API bridge policy rejects whitespace-padded tokens', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.BRIDGE_WORKER_ID = 'engineering-vm'; + env.BRIDGE_TOKEN = ' padded-development-bridge-token '; + env.PTC_MODE = 'replay'; + + expect(() => validateApiBridgePolicy()).toThrow( + 'must not contain surrounding whitespace', + ); + }); + test('remote bridge requires a positive finite job timeout', () => { env.SANDBOX_BACKEND = 'remote-bridge'; env.BRIDGE_WORKER_ID = 'engineering-vm'; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 357a2977..0bcf02c1 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -64,6 +64,11 @@ export function validateApiBridgePolicy(): void { 'CODEAPI_BRIDGE_WORKER_ID must match the bridge worker ID format', ); } + if (env.BRIDGE_TOKEN !== env.BRIDGE_TOKEN.trim()) { + throw new SecureStartupConfigError( + 'CODEAPI_BRIDGE_TOKEN must not contain surrounding whitespace', + ); + } if (env.HARDENED_SANDBOX_MODE) { requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); } else { From fb8391ab6239556d56bda44473cd74c52fb826dc Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 01:17:46 -0400 Subject: [PATCH 27/29] fix: close bridge registration and deadline races --- packages/code/src/cli.test.ts | 23 +++++++++++ packages/code/src/cli.ts | 23 +++++++---- packages/code/src/protocol.test.ts | 41 ++++++++++++++++++- packages/code/src/protocol.ts | 27 +++++++++++++ service/src/bridge/router.ts | 18 +-------- service/src/bridge/store.test.ts | 63 ++++++++++++++++++++++++++++++ service/src/bridge/store.ts | 23 +++++++++-- 7 files changed, 190 insertions(+), 28 deletions(-) diff --git a/packages/code/src/cli.test.ts b/packages/code/src/cli.test.ts index 837e9e99..179bf813 100644 --- a/packages/code/src/cli.test.ts +++ b/packages/code/src/cli.test.ts @@ -24,3 +24,26 @@ test('CLI rejects an invalid worker ID before entering the run loop', () => { /LIBRECHAT_CODE_WORKER_ID must match the bridge worker ID format/, ); }); + +test('CLI rejects invalid advertised capabilities before registration', () => { + const result = spawnSync( + process.execPath, + [fileURLToPath(new URL('./cli.js', import.meta.url))], + { + encoding: 'utf8', + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'https://code.example/v1', + LIBRECHAT_CODE_WORKER_TOKEN: 'worker-secret', + LIBRECHAT_CODE_WORKER_ID: 'engineering-vm', + LIBRECHAT_CODE_SANDBOX_PROFILE: '', + }, + }, + ); + + assert.notEqual(result.status, 0); + assert.match( + result.stderr, + /LIBRECHAT_CODE_SANDBOX_PROFILE or LIBRECHAT_CODE_RUNTIMES is invalid/, + ); +}); diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index a7f60cd7..cb1b100d 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -1,7 +1,10 @@ #!/usr/bin/env node import { createHash } from 'node:crypto'; import { BridgeWorker } from './worker.js'; -import { isValidBridgeWorkerId } from './protocol.js'; +import { + isValidBridgeWorkerCapabilities, + isValidBridgeWorkerId, +} from './protocol.js'; function required(name: string): string { const value = process.env[name]?.trim(); @@ -39,17 +42,23 @@ if (!isValidBridgeWorkerId(workerId)) { 'LIBRECHAT_CODE_WORKER_ID must match the bridge worker ID format', ); } +const capabilities = { + statefulWorkspace, + sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? 'nsjail', + runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), + policyDigest: createHash('sha256').update(policy).digest('hex'), +}; +if (!isValidBridgeWorkerCapabilities(capabilities)) { + throw new Error( + 'LIBRECHAT_CODE_SANDBOX_PROFILE or LIBRECHAT_CODE_RUNTIMES is invalid', + ); +} const worker = new BridgeWorker({ codeApiUrl: required('LIBRECHAT_CODE_URL'), token: required('LIBRECHAT_CODE_WORKER_TOKEN'), workerId, 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'), - }, + capabilities, onError: (error) => { const message = error instanceof Error ? error.message : 'unknown bridge error'; process.stderr.write(`librechat-code: reconnecting after ${message}\n`); diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts index 38f859e2..61c50922 100644 --- a/packages/code/src/protocol.test.ts +++ b/packages/code/src/protocol.test.ts @@ -1,6 +1,10 @@ import assert from 'node:assert/strict'; import test from 'node:test'; -import { bridgeWorkerPath, isValidBridgeWorkerId } from './protocol.js'; +import { + bridgeWorkerPath, + isValidBridgeWorkerCapabilities, + isValidBridgeWorkerId, +} from './protocol.js'; test('bridgeWorkerPath encodes worker-controlled path segments', () => { assert.equal( @@ -15,3 +19,38 @@ test('bridge worker IDs reject path, whitespace, and oversized values', () => { assert.equal(isValidBridgeWorkerId('engineering vm'), false); assert.equal(isValidBridgeWorkerId('a'.repeat(129)), false); }); + +test('bridge worker capabilities enforce registration limits', () => { + const valid = { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + policyDigest: 'a'.repeat(64), + }; + assert.equal(isValidBridgeWorkerCapabilities(valid), true); + assert.equal( + isValidBridgeWorkerCapabilities({ ...valid, sandboxProfile: '' }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + sandboxProfile: 'a'.repeat(129), + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + runtimes: Array.from({ length: 33 }, () => 'bash'), + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + runtimes: ['a'.repeat(65)], + }), + false, + ); +}); diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index dd3f54b5..9a5d2ae0 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -1,5 +1,8 @@ export const BRIDGE_PROTOCOL_VERSION = 1 as const; export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; +export const BRIDGE_SANDBOX_PROFILE_MAX_LENGTH = 128; +export const BRIDGE_RUNTIME_MAX_COUNT = 32; +export const BRIDGE_RUNTIME_MAX_LENGTH = 64; export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; @@ -100,3 +103,27 @@ export function bridgeWorkerPath(workerId: string): string { export function isValidBridgeWorkerId(workerId: string): boolean { return BRIDGE_WORKER_ID_PATTERN.test(workerId); } + +export function isValidBridgeWorkerCapabilities( + value: unknown, +): value is BridgeWorkerCapabilities { + if (typeof value !== 'object' || value === null) return false; + const capabilities = value as Record; + return ( + typeof capabilities.statefulWorkspace === 'boolean' && + typeof capabilities.sandboxProfile === 'string' && + capabilities.sandboxProfile.trim().length > 0 && + capabilities.sandboxProfile.length <= BRIDGE_SANDBOX_PROFILE_MAX_LENGTH && + Array.isArray(capabilities.runtimes) && + capabilities.runtimes.length <= BRIDGE_RUNTIME_MAX_COUNT && + capabilities.runtimes.every( + (runtime) => + typeof runtime === 'string' && + runtime.length > 0 && + runtime.length <= BRIDGE_RUNTIME_MAX_LENGTH, + ) && + (capabilities.policyDigest === undefined || + (typeof capabilities.policyDigest === 'string' && + /^[a-f0-9]{64}$/.test(capabilities.policyDigest))) + ); +} diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 3f03a533..35960401 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -7,6 +7,7 @@ import type { CodeBridgeAssignment, CodeBridgeSettlement } from './store'; import { BRIDGE_PROTOCOL_VERSION, + isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, } from '../../../packages/code/src/protocol'; import { connection } from '../queue'; @@ -110,22 +111,7 @@ router.post( 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))) + !isValidBridgeWorkerCapabilities(registration.capabilities) ) { res.status(400).json({ error: 'Invalid bridge worker registration' }); return; diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index e28b32e3..a7a10335 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -1560,6 +1560,69 @@ describe('RedisBridgeStore', () => { ).toHaveLength(0); }); + test('atomically rejects a fulfillment committed after its deadline', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'late-fulfillment-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const deadlineAtMs = Date.now() + 250; + const completion = store.dispatch({ + workerId: 'late-fulfillment-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-late-fulfillment', + deadlineAtMs, + signal: new AbortController().signal, + }); + const assignment = await store.lease( + 'late-fulfillment-worker', + incarnationId, + 1_000, + ); + await store.acknowledgeLease( + 'late-fulfillment-worker', + incarnationId, + assignment?.assignmentId ?? '', + assignment?.generation ?? 0, + assignment?.leaseToken ?? '', + ); + redis.eval = (async (...args: Parameters) => { + if ( + String(args[0]).includes("local existing = redis.call('GET', KEYS[2])") + ) { + await new Promise((resolve) => + setTimeout(resolve, Math.max(0, deadlineAtMs - Date.now() + 25)), + ); + } + return redisEval(...args); + }) as Redis['eval']; + + await expect( + store.settle('late-fulfillment-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-late-fulfillment', + files: [], + }, + }), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + test('releases the worker lock when generation allocation fails', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 0b2a84e5..512b4278 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -105,6 +105,10 @@ function settlementKey(assignmentId: string): string { return `${PREFIX}:assignment:${assignmentId}:settlement`; } +function assignmentDeadlineKey(assignmentId: string): string { + return `${PREFIX}:assignment:${assignmentId}:deadline`; +} + function cancellationKey(assignmentId: string): string { return `${PREFIX}:assignment:${assignmentId}:cancelled`; } @@ -800,6 +804,7 @@ export class RedisBridgeStore { settlementKey(assignmentId), leaseClaimKey(workerId, assignment.incarnationId), leaseAckKey(workerId, assignment.incarnationId), + assignmentDeadlineKey(assignmentId), ]; if (assignment.runtimeSessionId !== undefined) { settlementKeys.push( @@ -813,10 +818,11 @@ export class RedisBridgeStore { ' return -1', 'end', 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return 0 end', - 'if #KEYS == 5 and redis.call(\'GET\', KEYS[5]) ~= ARGV[3] then return -2 end', + 'if #KEYS == 6 and redis.call(\'GET\', KEYS[6]) ~= ARGV[3] then return -2 end', + 'if ARGV[4] ~= "rejected" and redis.call(\'EXISTS\', KEYS[5]) == 0 then return -3 end', 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', 'if redis.call(\'GET\', KEYS[3]) == ARGV[3] then redis.call(\'DEL\', KEYS[3], KEYS[4]) end', - 'if #KEYS == 5 and ARGV[4] == \"rejected\" then redis.call(\'DEL\', KEYS[5]) end', + 'if #KEYS == 6 and ARGV[4] == \"rejected\" then redis.call(\'DEL\', KEYS[6]) end', 'return 1', ].join('\n'); const accepted = Number( @@ -846,6 +852,12 @@ export class RedisBridgeStore { 'Bridge workspace in-flight marker was lost before settlement', ); } + if (accepted === -3) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment expired before settlement was committed', + ); + } if (accepted !== 1 && accepted !== 2) { throw new BridgeStoreError( 'ASSIGNMENT_EXPIRED', @@ -1068,12 +1080,13 @@ export class RedisBridgeStore { ): Promise { const script = [ 'if redis.call(\'GET\', KEYS[1]) ~= ARGV[1] then return 0 end', - 'if #KEYS == 5 and redis.call(\'EXISTS\', KEYS[5]) == 1 then return -1 end', + 'if #KEYS == 6 and redis.call(\'EXISTS\', KEYS[6]) == 1 then return -1 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])', - 'if #KEYS == 5 then redis.call(\'SET\', KEYS[5], ARGV[4]) end', + 'redis.call(\'SET\', KEYS[5], "1", \"PX\", ARGV[6])', + 'if #KEYS == 6 then redis.call(\'SET\', KEYS[6], ARGV[4]) end', 'return 1', ].join('\n'); const keys = [ @@ -1081,6 +1094,7 @@ export class RedisBridgeStore { assignmentKey(assignment.assignmentId), queueKey(assignment.workerId, assignment.incarnationId), lockIncarnationKey(assignment.workerId), + assignmentDeadlineKey(assignment.assignmentId), ]; if (assignment.runtimeSessionId !== undefined) { keys.push( @@ -1099,6 +1113,7 @@ export class RedisBridgeStore { String(ttlSeconds), assignment.assignmentId, String(ttlSeconds * 1000), + String(Math.max(1, Date.parse(assignment.expiresAt) - Date.now())), ); if (Number(result) === -1) { throw new BridgeStoreError( From 957e89082807bd0bb916094c06d1ea479616da04 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 01:27:57 -0400 Subject: [PATCH 28/29] fix: anchor bridge lease freshness --- packages/code/src/worker.test.ts | 68 ++++++++++++++++++++++++++++++++ packages/code/src/worker.ts | 3 +- service/src/bridge/store.test.ts | 23 +++++------ service/src/bridge/store.ts | 4 +- 4 files changed, 84 insertions(+), 14 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 94945cfd..27cf18d4 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -337,6 +337,74 @@ test('worker refreshes its registration during a long assignment', async () => { assert.ok(registrations >= 2); }); +test('worker schedules registration freshness from request start', async () => { + let registrations = 0; + const fetchImpl: typeof fetch = async (input) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + registrations += 1; + if (registrations === 1) { + await new Promise((resolve) => setTimeout(resolve, 40)); + } + 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, 10)); + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + if (url.endsWith('/cancelled')) { + return new Response( + JSON.stringify({ protocolVersion: 1, cancelled: false }), + { 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/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + registrationTransportTimeoutMs: 100, + cancellationPollIntervalMs: 100, + fetchImpl, + }); + await worker.register(); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-registration-transit', + 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 continues cancellation polling after a stalled response', async () => { let cancellationAttempts = 0; let settlementAttempted = false; diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 4fc4a78c..7d76e26b 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -110,6 +110,7 @@ export class BridgeWorker { ), ); const timeout = setTimeout(abortRegistration, timeoutMs); + const registrationStartedAtMs = Date.now(); let registration: BridgeWorkerRegistrationResponse; try { registration = await this.request( @@ -132,7 +133,7 @@ export class BridgeWorker { ); } this.registrationTtlMs = registration.leaseTtlMs; - this.lastRegisteredAtMs = Date.now(); + this.lastRegisteredAtMs = registrationStartedAtMs; return registration; } diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index a7a10335..24637a20 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -1572,6 +1572,18 @@ describe('RedisBridgeStore', () => { }, }); const deadlineAtMs = Date.now() + 250; + redis.eval = (async (...args: Parameters) => { + const script = String(args[0]); + if (script.includes("redis.call('RPUSH', KEYS[3], ARGV[4])")) { + await new Promise((resolve) => setTimeout(resolve, 75)); + } + if (script.includes("local existing = redis.call('GET', KEYS[2])")) { + await new Promise((resolve) => + setTimeout(resolve, Math.max(0, deadlineAtMs - Date.now() + 25)), + ); + } + return redisEval(...args); + }) as Redis['eval']; const completion = store.dispatch({ workerId: 'late-fulfillment-worker', body: { language: 'bash' } as t.PayloadBody, @@ -1592,17 +1604,6 @@ describe('RedisBridgeStore', () => { assignment?.generation ?? 0, assignment?.leaseToken ?? '', ); - redis.eval = (async (...args: Parameters) => { - if ( - String(args[0]).includes("local existing = redis.call('GET', KEYS[2])") - ) { - await new Promise((resolve) => - setTimeout(resolve, Math.max(0, deadlineAtMs - Date.now() + 25)), - ); - } - return redisEval(...args); - }) as Redis['eval']; - await expect( store.settle('late-fulfillment-worker', assignment?.assignmentId ?? '', { protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 512b4278..9b2dce13 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -1085,7 +1085,7 @@ export class RedisBridgeStore { 'redis.call(\'RPUSH\', KEYS[3], ARGV[4])', 'redis.call(\'EXPIRE\', KEYS[3], ARGV[3])', 'redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])', - 'redis.call(\'SET\', KEYS[5], "1", \"PX\", ARGV[6])', + 'redis.call(\'SET\', KEYS[5], "1", \"PXAT\", ARGV[6])', 'if #KEYS == 6 then redis.call(\'SET\', KEYS[6], ARGV[4]) end', 'return 1', ].join('\n'); @@ -1113,7 +1113,7 @@ export class RedisBridgeStore { String(ttlSeconds), assignment.assignmentId, String(ttlSeconds * 1000), - String(Math.max(1, Date.parse(assignment.expiresAt) - Date.now())), + String(Date.parse(assignment.expiresAt)), ); if (Number(result) === -1) { throw new BridgeStoreError( From 8a2ca2cc2230632e3372cbc8d9eb3dce7171e3e5 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 01:35:48 -0400 Subject: [PATCH 29/29] fix: preserve bridge response status --- packages/code/src/worker.test.ts | 77 ++++++++++++++++++++++++++++++++ packages/code/src/worker.ts | 14 ++++-- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 27cf18d4..d633366e 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -615,6 +615,54 @@ test('worker surfaces a definite stateless settlement rejection directly', async ); }); +test('worker preserves status for a non-JSON settlement rejection', async () => { + let settlementAttempts = 0; + 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: async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + settlementAttempts += 1; + return new Response('assignment fenced', { + status: 409, + headers: { 'Content-Type': 'text/html' }, + }); + }, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'non-json-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(), + request: { body: { language: 'bash' }, headers: {} }, + }), + (error: unknown) => + error instanceof Error && + error.name === 'BridgeProtocolError' && + 'status' in error && + error.status === 409, + ); + assert.equal(settlementAttempts, 1); +}); + test('worker retries an ambiguous settlement before the deadline', async () => { let settlementAttempts = 0; const fetchImpl: typeof fetch = async (input) => { @@ -1545,6 +1593,35 @@ test('worker bounds a stalled registration below its lease TTL', async () => { await assert.rejects(worker.register(), { name: 'AbortError' }); }); +test('worker preserves status for a non-JSON registration rejection', async () => { + 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: async () => + new Response('unauthorized', { + status: 401, + headers: { 'Content-Type': 'text/html' }, + }), + }); + + await assert.rejects( + worker.register(), + (error: unknown) => + error instanceof Error && + error.name === 'BridgeProtocolError' && + 'status' in error && + error.status === 401, + ); +}); + test('worker uses the server-relative lease budget despite VM clock skew', async () => { let settlementAttempted = false; const worker = new BridgeWorker({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 7d76e26b..5c1afb95 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -752,13 +752,21 @@ export class BridgeWorker { body: JSON.stringify(body), signal, }); - const payload = (await response.json()) as object; + let payload: unknown; + try { + payload = await response.json(); + } catch (error) { + if (response.ok) throw error; + payload = {}; + } if (!response.ok) { + const errorPayload = + typeof payload === 'object' && payload !== null ? payload : {}; throw new BridgeProtocolError( - errorMessage(payload) ?? + errorMessage(errorPayload) ?? `Bridge request failed with HTTP ${response.status}`, response.status, - errorCode(payload), + errorCode(errorPayload), ); } return payload as T;