diff --git a/.env.example b/.env.example index 54a2971a..98c97260 100644 --- a/.env.example +++ b/.env.example @@ -24,6 +24,14 @@ 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 +# CODEAPI_BRIDGE_AUTH_MODE=paired + # Service Configuration PYTHON_CONCURRENCY=5 OTHER_CONCURRENCY=15 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 94e22864..7ba2ef1c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,6 +34,9 @@ jobs: - name: Sandbox-runner liveness checks run: tests/sandbox_runner_healthcheck.sh + - name: Bridge pairing rollout safety + run: tests/bridge_pairing_rollout.sh + - name: Validate sandbox Dockerfiles run: | docker buildx build --check -f api/Dockerfile . 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..384e013a 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 @@ -35,7 +37,9 @@ Set `CODEAPI_EXECUTION_PROFILE` consistently on an API deployment and its workers. The default profile keeps the existing `python-queue` and `other-queue`; the stateful profile uses `stateful-python-queue` and `stateful-other-queue`. This allows both deployments to share Redis without -cross-consuming jobs. +cross-consuming jobs. The `remote-bridge` backend additionally uses +`remote-bridge-python-queue` and `remote-bridge-other-queue`, fencing attached +worker jobs from Lambda consumers during rolling deployments. An existing Lambda MicroVM deployment upgraded from a pre-profile release may leave `CODEAPI_EXECUTION_PROFILE` unset for its first binary rollout. An @@ -65,6 +69,18 @@ Two modes are supported: - **NsJail mode** (`kvmEnabled: false`): Direct NsJail sandboxing with Linux namespaces and cgroups - **MicroVM mode** (`kvmEnabled: true`): libkrun microVM with its own kernel, NsJail runs inside the guest +## Remote stateful environments + +The `remote-bridge` backend keeps the Code API as the policy and queue boundary +while moving execution to a sandbox on an operator-selected VM. The worker only +makes outbound authenticated requests, so the VM does not need a public ingress +port. Assignments carry a deadline, a single-active-worker lock, a monotonically +increasing generation, and a one-time lease token to fence stale workers. + +See [Remote Code Bridge](docs/remote-bridge/README.md) for deployment and threat +model details. The worker protocol and CLI live in the provider-neutral +[`@librechat/code`](packages/code/README.md) package. + ## Security disclaimer This service exists to run arbitrary, untrusted code — treat every diff --git a/api/Dockerfile b/api/Dockerfile index f8d6713c..3526e98e 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -29,8 +29,10 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \ RUN make -j$(nproc) COPY api/src/spec-guard.c /tmp/spec-guard.c +COPY docker/rootfs-setup.c /tmp/rootfs-setup.c RUN gcc -O2 -static -o /usr/local/bin/spec-guard /tmp/spec-guard.c \ - && chmod 0111 /usr/local/bin/spec-guard + && gcc -O2 -static -o /usr/local/bin/sandbox-rootfs-setup /tmp/rootfs-setup.c \ + && chmod 0111 /usr/local/bin/spec-guard /usr/local/bin/sandbox-rootfs-setup # ============================================================================ # Stage 1b: Build language runtime packages (only consumed by sandbox-runner-baked) @@ -212,6 +214,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ && dnf clean all COPY --from=launcher-builder /launcher/target/release/sandbox-launcher /usr/local/bin/launcher +COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-setup COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh COPY docker/start-direct-sandbox.sh /usr/local/bin/start-direct-sandbox.sh diff --git a/api/src/entrypoint.sh b/api/src/entrypoint.sh index b532a67c..4fce5ba4 100755 --- a/api/src/entrypoint.sh +++ b/api/src/entrypoint.sh @@ -171,28 +171,33 @@ fi chmod 777 "$SMOKE_DIR" fi SMOKE_LOG=$(mktemp) +SMOKE_STDERR=$(mktemp) NSJAIL_CGROUP_ARGS=() if [ "$SANDBOX_USE_CGROUPV2" = "true" ]; then NSJAIL_CGROUP_ARGS=(--use_cgroupv2) fi -if timeout 10 /usr/sbin/nsjail --config "${NSJAIL_CONFIG:-/sandbox_api/config/sandbox.cfg}" \ +if timeout 10 "${NSJAIL_PATH:-/usr/sbin/nsjail}" --config "${NSJAIL_CONFIG:-/sandbox_api/config/sandbox.cfg}" \ "${NSJAIL_CGROUP_ARGS[@]}" --log "$SMOKE_LOG" \ --user "65534:${SMOKE_OUTSIDE_UID}:1" --group "65534:${SMOKE_OUTSIDE_GID}:1" \ -s /usr/bin:/bin -s /usr/lib:/lib -s /usr/lib64:/lib64 \ -B "$SMOKE_DIR:/mnt/data" \ - -- /bin/sh -c 'printf "%s\n" sandbox_ok > /mnt/data/smoke.txt && test "$(cat /mnt/data/smoke.txt)" = sandbox_ok' > /dev/null 2>&1; then + -- /bin/sh -c 'printf "%s\n" sandbox_ok > /mnt/data/smoke.txt && test "$(cat /mnt/data/smoke.txt)" = sandbox_ok' > /dev/null 2>"$SMOKE_STDERR"; then echo "NsJail smoke test passed" else echo "FATAL: NsJail smoke test failed — sandbox cannot start" echo "NsJail log output:" cat "$SMOKE_LOG" 2>/dev/null || true + echo "NsJail stderr:" + cat "$SMOKE_STDERR" 2>/dev/null || true rm -f "$SMOKE_LOG" + rm -f "$SMOKE_STDERR" rm -rf "$SMOKE_DIR" exit 1 fi rm -f "$SMOKE_LOG" +rm -f "$SMOKE_STDERR" rm -rf "$SMOKE_DIR" echo "Starting sandbox API server..." diff --git a/docker/Dockerfile.worker-sandbox b/docker/Dockerfile.worker-sandbox index cd3edd3d..c18eab82 100644 --- a/docker/Dockerfile.worker-sandbox +++ b/docker/Dockerfile.worker-sandbox @@ -42,8 +42,10 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \ RUN make -j$(nproc) COPY api/src/spec-guard.c /tmp/spec-guard.c +COPY docker/rootfs-setup.c /tmp/rootfs-setup.c RUN gcc -O2 -static -o /usr/local/bin/spec-guard /tmp/spec-guard.c \ - && chmod 0111 /usr/local/bin/spec-guard + && gcc -O2 -static -o /usr/local/bin/sandbox-rootfs-setup /tmp/rootfs-setup.c \ + && chmod 0111 /usr/local/bin/spec-guard /usr/local/bin/sandbox-rootfs-setup # ============================================================================ # Stage 1b: Build language runtime packages for the baked KVM root disk @@ -83,6 +85,7 @@ WORKDIR /app COPY service/package.json service/bun.lock ./ RUN bun install --frozen-lockfile COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ RUN bun build ./src/worker-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' @@ -231,6 +234,7 @@ ENV PATH="/root/.bun/bin:${PATH}" # --- Launcher (runs on host, boots microVM) --- COPY --from=launcher-builder /launcher/target/release/sandbox-launcher /usr/local/bin/launcher +COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-setup # --- Launcher entrypoint (DNS resolution + socat relay before VM boot) --- COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh diff --git a/docker/rootfs-setup.c b/docker/rootfs-setup.c new file mode 100644 index 00000000..b0b817cf --- /dev/null +++ b/docker/rootfs-setup.c @@ -0,0 +1,92 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +static int bind_mount(const char *source, const char *target, int read_only) { + if (mount(source, target, NULL, MS_BIND | MS_REC, NULL) != 0) { + fprintf(stderr, "bind %s -> %s failed: %s\n", source, target, strerror(errno)); + return -1; + } + + if (read_only && + mount(NULL, target, NULL, MS_BIND | MS_REMOUNT | MS_RDONLY, NULL) != 0) { + fprintf(stderr, "read-only remount of %s failed: %s\n", target, strerror(errno)); + return -1; + } + + return 0; +} + +static int bind_rootfs_path(const char *rootfs, const char *path) { + char source[PATH_MAX]; + int written = snprintf(source, sizeof(source), "%s%s", rootfs, path); + if (written < 0 || (size_t)written >= sizeof(source)) { + fprintf(stderr, "rootfs path is too long: %s%s\n", rootfs, path); + return -1; + } + + return bind_mount(source, path, 1); +} + +int main(int argc, char **argv) { + if (argc < 2) { + fprintf(stderr, "usage: sandbox-rootfs-setup ROOTFS [COMMAND ...]\n"); + return 2; + } + + const char *rootfs = argv[1]; + if (rootfs[0] != '/') { + fprintf(stderr, "rootfs must be an absolute path\n"); + return 2; + } + + if (mount(NULL, "/", NULL, MS_REC | MS_PRIVATE, NULL) != 0) { + fprintf(stderr, "making the mount namespace private failed: %s\n", strerror(errno)); + return 1; + } + + if ((mkdir("/sandbox_api", 0755) != 0 && errno != EEXIST) || + (mkdir("/pkgs", 0755) != 0 && errno != EEXIST)) { + fprintf(stderr, "creating rootfs mount targets failed: %s\n", strerror(errno)); + return 1; + } + + /* + * Keep this process statically linked: the final /usr mount replaces + * the Fedora launcher's dynamic userspace with the Debian sandbox rootfs. + * A shell cannot safely perform this sequence because its next command may + * try to load a host binary against guest libraries (or vice versa). + */ + const char *paths[] = {"/sandbox_api", "/pkgs"}; + for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); i++) { + if (bind_rootfs_path(rootfs, paths[i]) != 0) { + return 1; + } + } + + if (access("/host-packages", F_OK) == 0 && + bind_mount("/host-packages", "/pkgs", 0) != 0) { + fprintf(stderr, "warning: sandbox will run without host packages\n"); + } + + /* Bind all guest userspace last, then immediately enter it. */ + if (bind_rootfs_path(rootfs, "/usr") != 0) { + return 1; + } + + setenv("PATH", "/root/.bun/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", 1); + setenv("LD_LIBRARY_PATH", "/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu", 1); + + setenv("NSJAIL_PATH", "/usr/sbin/nsjail", 1); + + char *default_argv[] = {"/sandbox_api/entrypoint.sh", NULL}; + char **command_argv = argc > 2 ? &argv[2] : default_argv; + execv(command_argv[0], command_argv); + fprintf(stderr, "starting sandbox entrypoint failed: %s\n", strerror(errno)); + return 1; +} diff --git a/docker/start-direct-sandbox.sh b/docker/start-direct-sandbox.sh index a171a6df..bf7c9805 100644 --- a/docker/start-direct-sandbox.sh +++ b/docker/start-direct-sandbox.sh @@ -44,35 +44,4 @@ else fi export SANDBOX_ROOTFS="$ROOTFS" - -exec unshare --mount bash -c ' - ROOTFS="${SANDBOX_ROOTFS:-/sandbox-rootfs}" - - mount -o bind,ro "$ROOTFS/usr/sbin" /usr/sbin || { echo "FATAL: cannot bind /usr/sbin"; exit 1; } - mount -o bind,ro "$ROOTFS/usr/lib" /usr/lib || { echo "FATAL: cannot bind /usr/lib"; exit 1; } - - if [ -d "$ROOTFS/usr/lib64" ] && ! [ -L "$ROOTFS/usr/lib64" ]; then - mount -o bind,ro "$ROOTFS/usr/lib64" /usr/lib64 2>/dev/null || \ - echo "[sandbox] WARNING: could not bind /usr/lib64 - sandboxed binaries may fail to exec" - fi - - mount -o bind,ro "$ROOTFS/usr/local" /usr/local || { echo "FATAL: cannot bind /usr/local"; exit 1; } - mount -o bind,ro "$ROOTFS/sandbox_api" /sandbox_api || { echo "FATAL: cannot bind /sandbox_api"; exit 1; } - mount -o bind,ro "$ROOTFS/pkgs" /pkgs || { echo "FATAL: cannot bind /pkgs"; exit 1; } - - if [ -d /host-packages ]; then - mount --bind /host-packages /pkgs 2>/dev/null || \ - echo "WARNING: could not bind /host-packages - sandbox will run without packages" - fi - - mount -o bind,ro "$ROOTFS/usr/bin" /usr/bin || { echo "FATAL: cannot bind /usr/bin"; exit 1; } - - multiarch_libdir=$(find /usr/lib -maxdepth 1 -type d -name "*-linux-gnu" -print -quit) - if [ -n "$multiarch_libdir" ]; then - export LD_LIBRARY_PATH="$multiarch_libdir${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" - fi - - export PATH="/root/.bun/bin:$PATH" - - exec /sandbox_api/entrypoint.sh -' +exec unshare --mount /sandbox-rootfs-setup "$ROOTFS" diff --git a/docs/adr/001-stateful-code-environments.md b/docs/adr/001-stateful-code-environments.md new file mode 100644 index 00000000..8b9830df --- /dev/null +++ b/docs/adr/001-stateful-code-environments.md @@ -0,0 +1,86 @@ +# ADR 001: Stateful code environments use an outbound Code API bridge + +- Status: Accepted for alpha +- Date: 2026-08-30 + +## Context + +LibreChat needs coding agents to reuse a workspace across conversation turns +while allowing the environment owner to choose the VM. Internet-facing +LibreChat instances cannot safely require inbound access to that VM, forward +end-user tokens to it, or treat an MCP connection as a sandbox boundary. + +The first alpha demonstrated a stable runtime-session ID, a single fenced +worker lease, and workspace persistence across turns. Its static shared worker +token was sufficient to prove execution flow but is not an acceptable hardened +enrollment mechanism. + +## Decision + +The product concept is a **stateful code environment**. Code API remains its +broker and policy boundary, and `remote-bridge` is a Code API sandbox backend. +The `@librechat/code` worker connects outbound from the chosen VM and forwards +assignments only to a loopback or private sandbox endpoint. + +Hardened workers enroll through a one-time pairing code: + +1. An administrator creates a code scoped to the configured worker ID. +2. The CLI generates an Ed25519 keypair locally and redeems the code with only + its public key. +3. Code API returns a fifteen-minute credential bound to that public key. +4. Every worker request signs the method, path, body digest, timestamp, nonce, + and credential. +5. Code API rejects stale timestamps and replayed nonces and supports rotation + and immediate revocation. + +Static bearer authentication remains a non-hardened compatibility mode. + +## Ownership and state + +The alpha environment is deployment/operator owned and configured with one +worker ID. A future LibreChat control plane may persist deployment-, tenant-, +or user-owned environment records and issue the same pairing operation through +RBAC-protected APIs without changing the worker execution protocol. + +Workspace state belongs to the stable runtime session, not to a transient +assignment lease. For `remote-bridge`, that state currently survives turns on +the same worker and backing disk. It is not yet checkpointed or portable across +worker replacement; the UI and operator documentation must not imply otherwise. + +## Security invariants + +- The VM requires no inbound internet listener. +- Code API, not the worker, authenticates LibreChat users and normalizes work. +- A stolen short-lived credential is insufficient without the worker private + key; a stolen private key is insufficient after credential expiry or + revocation. +- Pairing codes and credentials are stored by digest where lookup permits. +- One configured worker has at most one active fenced assignment. +- Sandbox isolation and default-deny egress remain mandatory; pairing secures + the transport identity but does not make the host a sandbox. +- A compromised worker can lie about advertised capabilities. Capability + labels and policy digests are audit signals until enforcement is coupled to + an attested sandbox or trusted host policy. + +## Consequences + +- `@librechat/code` owns the provider-neutral protocol, identity handling, and + worker CLI; Code API owns enrollment, scheduling, and execution policy. +- LibreChat owns environment persistence, ownership, RBAC, and user experience. +- The Agents SDK keeps only its adapter until a second concrete consumer proves + which coding-tool abstractions are genuinely provider neutral. +- MCP may expose environment operations later, but it is not the worker + transport or isolation boundary. +- Multi-worker directories, checkpoint/restore, owner-scoped quotas, and + enforced network capability profiles remain follow-up decisions. + +## Alternatives rejected + +- **Inbound SSH/HTTP to the VM:** expands attack surface and complicates NAT and + firewall operation. +- **MCP as the worker protocol:** conflates tool discovery with leases, + cancellation, fencing, and sandbox policy. +- **Put the runtime in the Agents SDK:** couples provider-neutral execution to + one agent integration and makes non-agent consumers depend on agent internals. +- **Long-lived shared bearer token:** easy to bootstrap, but replayable and not + bound to a worker-held key. diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md new file mode 100644 index 00000000..8d10f024 --- /dev/null +++ b/docs/remote-bridge/README.md @@ -0,0 +1,177 @@ +# 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 +identity 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= +CODEAPI_BRIDGE_AUTH_MODE=paired +``` + +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. + +To attach multiple principal-owned workers to one Code API deployment, enable +dynamic routing. A compatibility default worker is optional in this mode: + +```dotenv +CODEAPI_BRIDGE_DYNAMIC_WORKERS=true +CODEAPI_BRIDGE_AUTH_MODE=paired +# CODEAPI_BRIDGE_WORKER_ID=my-default-vm +``` + +Dynamic routing is accepted only with paired authentication. LibreChat signs +the selected worker into the short-lived Code API JWT as `code_worker_id`. +`X-LibreChat-Code-Worker-ID` remains the transport header, but Code API accepts +it only when it exactly matches that authenticated claim. The resolved worker +is persisted across the queue and programmatic replay boundaries, and Code API +requires both its stored tenant binding and registered worker credential before +creating a lease. + +Create a single-use pairing code with the administrator secret: + +```bash +curl -fsS https://code.example.com/v1/bridge/pairings \ + -H "Authorization: Bearer $CODEAPI_BRIDGE_TOKEN" \ + -H 'Content-Type: application/json' \ + --data '{"workerId":"my-vm"}' +``` + +With dynamic routing enabled, the trusted control plane must bind each pairing +to one tenant and generic principal. Code API treats the principal as lifecycle +and audit metadata; LibreChat remains responsible for resolving user, role, and +group membership before selecting the worker: + +```bash +curl -fsS https://code.example.com/v1/bridge/pairings \ + -H "Authorization: Bearer $CODEAPI_BRIDGE_TOKEN" \ + -H 'Content-Type: application/json' \ + --data '{ + "workerId":"user-vm", + "binding":{ + "tenantId":"tenant-1", + "principal":{"type":"user","id":"user-1"} + } + }' +``` + +Principal types are `deployment`, `tenant`, `user`, `role`, and `group`. +Pairing and registration bodies from the VM cannot replace the server-issued +binding, and credential rotation preserves it. + +Redeem the returned code on the VM using +[`@librechat/code`](../../packages/code/README.md). The CLI generates its key +locally, proves possession on every request, and rotates its short-lived +credential before expiry. `CODEAPI_BRIDGE_AUTH_MODE=static` remains available +for non-hardened development compatibility only. + +Stateful deployments must also set `LIBRECHAT_CODE_STATEFUL_WORKSPACE=true` +and route the CLI's `{runtimeSessionId}` endpoint template to an isolated, +persistent local runner per session. A single sandbox endpoint is stateless and +is rejected for runtime-session assignments. + +## LibreChat configuration + +Expose the Code API deployment as an environment under the Agents endpoint: + +```yaml +endpoints: + agents: + statefulCodeSessions: + environments: + - id: my-vm + name: My VM + type: attached + baseURL: https://code.example.com/v1 + default: true +``` + +Agents may select this environment with `code_environment_id: my-vm`. +LibreChat derives a stable per-conversation runtime session ID, so commands in +later turns reuse the same workspace. Attached environments deliberately skip +background prewarming: the single worker lease is reserved for explicit user +execution. + +## Lifecycle and fencing + +- Registration is ephemeral in Redis and must be refreshed by the worker. +- Pairing codes are stored hashed, expire after ten minutes, and are consumed + atomically on their first redemption attempt. +- Worker credentials expire after fifteen minutes and are bound to an Ed25519 + public key. Exact-request signatures include the HTTP method, path, body + digest, timestamp, nonce, and credential. +- Accepted proof nonces cannot be replayed, credentials rotate before expiry, + and an administrator can revoke the active worker identity immediately. +- Assignment leases bind to a stable paired identity rather than an individual + short-lived credential. Rotation preserves that identity; pairing again + replaces it and fences work queued for the previous owner. +- Remote bridge deployments use backend-specific BullMQ queues and serialize + the expected backend on every new job, preventing Lambda or HTTP consumers + from accepting attached-worker executions. +- Code API permits one active assignment per worker. +- Dynamic workers are fenced to their server-issued tenant before assignment. +- 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. +- 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. +- 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; 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. +- 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. + +## 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 +administrator token immediately. Pairing secures worker transport identity; it +cannot attest that a compromised VM truthfully reports or enforces its sandbox +capabilities. + +LibreChat's owner-scoped environment registry can issue these principal-bound +pairings without changing the worker execution protocol or moving code tools +into the Agents SDK. diff --git a/helm/codeapi/README.md b/helm/codeapi/README.md index 9b3efbb0..fcdef01d 100644 --- a/helm/codeapi/README.md +++ b/helm/codeapi/README.md @@ -53,6 +53,32 @@ platform rather than templated here: external ingress/service mesh, KEDA-style queue-depth autoscaling, and cloud-IAM secret delivery (the env hooks below cover all of them). +**Pairing-fence rollbacks.** Do not use a direct `helm rollback` from a chart +revision containing the bridge pairing fence to an older revision. Helm runs +rollback hooks from the target revision, so a pre-fence target cannot stop its +own old and new API replicas from overlapping. Use the chart's fail-closed +helper instead: + +```bash +helm/codeapi/scripts/safe-pairing-rollback.sh RELEASE REVISION NAMESPACE +``` + +The helper records an out-of-band rollback epoch, deletes the API HPA, scales +the live fenced API deployment to zero, verifies that the Deployment and every +matching pod have converged to zero, and only then invokes `helm rollback`. +When a fenced revision is deployed again, the epoch forces one fresh cleanup of +legacy pairing codes even if the original migration window has expired. This +causes an API outage by design. If rollback fails, the helper repeats the drain +after re-discovering every API Deployment and explicitly deletes any remaining +API pods, so a partially applied rollback cannot leave a mixed-version API +running. The operator running it needs permission to read/scale Deployments, +delete HPAs and pods, and create or update the rollback ConfigMap. +Pass the intended cluster context to both `kubectl` and `helm` before invoking +the helper; it rejects forwarded kubeconfig, context, identity, API-server, and +namespace flags and Helm-specific target environment overrides so the drain and +rollback cannot target different clusters. Termination signals during Helm +also trigger a final recovery drain before the helper exits. + **Execution profile.** By default this chart leaves `CODEAPI_EXECUTION_PROFILE` unset. Its bundled HTTP/stateless configuration is inferred as the AWS-free `default` profile and retains the existing diff --git a/helm/codeapi/scripts/safe-pairing-rollback.sh b/helm/codeapi/scripts/safe-pairing-rollback.sh new file mode 100755 index 00000000..c70bddc7 --- /dev/null +++ b/helm/codeapi/scripts/safe-pairing-rollback.sh @@ -0,0 +1,190 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + echo "usage: $0 RELEASE REVISION [NAMESPACE] [helm rollback flags...]" >&2 + exit 64 +} + +release=${1:-} +revision=${2:-} +namespace=${3:-default} +if [[ -z "$release" || ! "$revision" =~ ^[1-9][0-9]*$ ]]; then + usage +fi +shift $(( $# >= 3 ? 3 : $# )) + +if [[ ! "$release" =~ ^[a-z0-9]([-a-z0-9]*[a-z0-9])?$ ]]; then + echo "invalid Helm release name: $release" >&2 + exit 64 +fi +if [[ ! "$namespace" =~ ^[a-z0-9]([-a-z0-9]*[a-z0-9])?$ ]]; then + echo "invalid Kubernetes namespace: $namespace" >&2 + exit 64 +fi +for flag in "$@"; do + case "$flag" in + -n|-n?*|--namespace|--namespace=*|--kube-context|--kube-context=*|\ + --kubeconfig|--kubeconfig=*|--kube-apiserver|--kube-apiserver=*|\ + --kube-ca-file|--kube-ca-file=*|--kube-token|--kube-token=*|\ + --kube-tls-server-name|--kube-tls-server-name=*|\ + --kube-as-user|--kube-as-user=*|--kube-as-group|--kube-as-group=*|\ + --kube-insecure-skip-tls-verify|--kube-insecure-skip-tls-verify=*) + echo "refusing target-changing Helm rollback flag: $flag" >&2 + exit 64 + ;; + esac +done +for variable in \ + HELM_KUBEAPISERVER \ + HELM_KUBEASGROUPS \ + HELM_KUBEASUSER \ + HELM_KUBECAFILE \ + HELM_KUBECONTEXT \ + HELM_KUBEINSECURE_SKIP_TLS_VERIFY \ + HELM_KUBETLS_SERVER_NAME \ + HELM_KUBETOKEN \ + HELM_NAMESPACE; do + if [[ -n ${!variable:-} ]]; then + echo "refusing Helm target override from environment: $variable" >&2 + exit 64 + fi +done + +timeout=${CODEAPI_ROLLBACK_TIMEOUT:-10m} +selector="app.kubernetes.io/instance=${release},app.kubernetes.io/component=api" + +discover_api_deployments() { + local output + output=$(kubectl --namespace "$namespace" get deployment \ + --selector "$selector" --output name) || return + deployments=() + if [[ -n "$output" ]]; then + mapfile -t deployments <<< "$output" + fi +} + +list_api_pods() { + local output + output=$(kubectl --namespace "$namespace" get pod \ + --selector "$selector" --output name) || return + pods=() + if [[ -n "$output" ]]; then + mapfile -t pods <<< "$output" + fi +} + +discover_api_deployments +if (( ${#deployments[@]} != 1 )); then + echo "expected exactly one Code API deployment for $selector" >&2 + exit 1 +fi +deployment=${deployments[0]} + +fence=$(kubectl --namespace "$namespace" get "$deployment" \ + --output 'jsonpath={.spec.template.metadata.annotations.codeapi\.librechat\.ai/pairing-fence-version}') +if [[ -z "$fence" ]]; then + echo "refusing rollback: the live API deployment has no pairing fence" >&2 + exit 1 +fi + +deployment_name=${deployment#*/} +rollback_config_map=${deployment_name%-api}-pairing-rollback +rollback_epoch="$(date +%s)-${RANDOM}-${RANDOM}" + +echo "Recording pairing rollback epoch $rollback_epoch..." >&2 +kubectl --namespace "$namespace" create configmap "$rollback_config_map" \ + --from-literal="epoch=$rollback_epoch" --dry-run=client --output yaml | \ + kubectl --namespace "$namespace" apply --filename - + +drain_api() { + local pod_action=${1:-wait} + local replica_state desired current ready available updated + + # Helm may have partially installed a target with a different fullname. + # Resolve every matching API Deployment on each drain attempt. + discover_api_deployments + if (( ${#deployments[@]} == 0 )) && [[ "$pod_action" != delete ]]; then + echo "refusing rollback: no API deployment matched $selector" >&2 + return 1 + fi + + echo "Deleting API autoscalers before the rollback fence is lowered..." >&2 + kubectl --namespace "$namespace" delete horizontalpodautoscaler \ + --selector "$selector" --ignore-not-found --wait=true + + echo "Scaling the fenced API deployment to zero..." >&2 + for deployment in "${deployments[@]}"; do + kubectl --namespace "$namespace" scale "$deployment" --replicas=0 + kubectl --namespace "$namespace" rollout status "$deployment" \ + --timeout "$timeout" + done + + list_api_pods + if (( ${#pods[@]} > 0 )); then + if [[ "$pod_action" == delete ]]; then + kubectl --namespace "$namespace" delete pod \ + --selector "$selector" --wait=true --timeout "$timeout" + else + kubectl --namespace "$namespace" wait "${pods[@]}" \ + --for=delete --timeout "$timeout" + fi + fi + + # Relist immediately before Helm can lower the fence. This catches a new + # matching pod that appeared after the first snapshot. + discover_api_deployments + for deployment in "${deployments[@]}"; do + replica_state=$(kubectl --namespace "$namespace" get "$deployment" \ + --output 'jsonpath={.spec.replicas},{.status.replicas},{.status.readyReplicas},{.status.availableReplicas},{.status.updatedReplicas}') || return + IFS=, read -r desired current ready available updated <<< "$replica_state" + if [[ ${desired:-0} != 0 || ${current:-0} != 0 || ${ready:-0} != 0 || + ${available:-0} != 0 || ${updated:-0} != 0 ]]; then + echo "refusing rollback: API deployment did not converge to zero replicas" >&2 + return 1 + fi + done + list_api_pods + if (( ${#pods[@]} > 0 )); then + echo "refusing rollback: API pods appeared after the drain" >&2 + return 1 + fi +} + +drain_api wait + +echo "All fenced API pods are gone; starting Helm rollback..." >&2 +rollback_pid= +recover_interrupted_rollback() { + local exit_status=$1 + trap - HUP INT TERM + if [[ -n "$rollback_pid" ]]; then + kill -TERM "$rollback_pid" 2>/dev/null || true + wait "$rollback_pid" 2>/dev/null || true + fi + echo "Helm rollback interrupted; restoring the fail-closed API drain..." >&2 + set -e + drain_api delete + exit "$exit_status" +} +trap 'recover_interrupted_rollback 129' HUP +trap 'recover_interrupted_rollback 130' INT +trap 'recover_interrupted_rollback 143' TERM + +helm rollback "$release" "$revision" \ + --namespace "$namespace" --wait --wait-for-jobs --timeout "$timeout" "$@" & +rollback_pid=$! +set +e +wait "$rollback_pid" +rollback_status=$? +set -e +rollback_pid= +trap - HUP INT TERM + +if (( rollback_status == 0 )); then + exit 0 +else + echo "Helm rollback failed; restoring the fail-closed API drain..." >&2 + drain_api delete + exit "$rollback_status" +fi diff --git a/helm/codeapi/templates/api-deployment.yaml b/helm/codeapi/templates/api-deployment.yaml index bf5f91e0..69985db4 100644 --- a/helm/codeapi/templates/api-deployment.yaml +++ b/helm/codeapi/templates/api-deployment.yaml @@ -17,11 +17,15 @@ spec: {{- if not .Values.api.autoscaling.enabled }} replicas: {{ .Values.api.replicaCount }} {{- end }} + strategy: + {{- toYaml .Values.api.strategy | nindent 4 }} selector: matchLabels: {{- include "codeapi.api.selectorLabels" . | nindent 6 }} template: metadata: + annotations: + codeapi.librechat.ai/pairing-fence-version: "1" labels: {{- include "codeapi.api.selectorLabels" . | nindent 8 }} spec: @@ -55,6 +59,12 @@ spec: secretKeyRef: name: {{ include "codeapi.fullname" . }}-secrets key: redis-password + - name: CODEAPI_BRIDGE_PAIRING_ROLLBACK_EPOCH + valueFrom: + configMapKeyRef: + name: {{ include "codeapi.fullname" . }}-pairing-rollback + key: epoch + optional: true # Service URLs - name: FILE_SERVER_URL value: "http://{{ include "codeapi.fullname" . }}-file-server:{{ .Values.fileServer.service.port }}" diff --git a/helm/codeapi/values.yaml b/helm/codeapi/values.yaml index 75abaaf7..ee997ce8 100644 --- a/helm/codeapi/values.yaml +++ b/helm/codeapi/values.yaml @@ -68,10 +68,20 @@ api: enabled: true replicaCount: 2 # Start with 2 API pods + # Pairing revocation relies on every serving replica honoring the Redis + # generation fence. Recreate prevents a pre-fence binary from redeeming an + # already-revoked code during the first rollout of paired bridge workers. + # Roll back to pre-fence revisions only with scripts/safe-pairing-rollback.sh. + strategy: + type: Recreate + rollingUpdate: null + image: repository: codeapi-api tag: latest - pullPolicy: IfNotPresent # Use Always in production + # Recreate is a security fence only if replacement pods cannot reuse a + # cached pre-fence image behind the mutable default tag. + pullPolicy: Always # Resource limits resources: diff --git a/launcher/Dockerfile b/launcher/Dockerfile index 0d252c28..1a077e12 100644 --- a/launcher/Dockerfile +++ b/launcher/Dockerfile @@ -33,8 +33,10 @@ RUN git clone -b master --single-branch https://github.com/google/nsjail.git . \ RUN make -j$(nproc) COPY api/src/spec-guard.c /tmp/spec-guard.c +COPY docker/rootfs-setup.c /tmp/rootfs-setup.c RUN gcc -O2 -static -o /usr/local/bin/spec-guard /tmp/spec-guard.c \ - && chmod 0111 /usr/local/bin/spec-guard + && gcc -O2 -static -o /usr/local/bin/sandbox-rootfs-setup /tmp/rootfs-setup.c \ + && chmod 0111 /usr/local/bin/spec-guard /usr/local/bin/sandbox-rootfs-setup FROM oven/bun:1.3.14-debian AS sandbox-build @@ -126,6 +128,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ && dnf clean all COPY --from=launcher-builder /launcher/target/release/sandbox-launcher /usr/local/bin/launcher +COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-setup COPY --from=sandbox-build / /sandbox-rootfs/ @@ -136,6 +139,6 @@ RUN mkdir -p /host-packages COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh COPY docker/start-direct-sandbox.sh /usr/local/bin/start-direct-sandbox.sh COPY docker/sandbox-entrypoint.sh /usr/local/bin/entrypoint.sh -RUN chmod +x /usr/local/bin/launcher-entrypoint.sh /usr/local/bin/start-direct-sandbox.sh /usr/local/bin/entrypoint.sh +RUN chmod +x /sandbox-rootfs-setup /usr/local/bin/launcher-entrypoint.sh /usr/local/bin/start-direct-sandbox.sh /usr/local/bin/entrypoint.sh ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] diff --git a/packages/code/Dockerfile b/packages/code/Dockerfile new file mode 100644 index 00000000..21fb09bd --- /dev/null +++ b/packages/code/Dockerfile @@ -0,0 +1,15 @@ +FROM node:24-alpine AS build +WORKDIR /app +COPY package.json package-lock.json tsconfig.json ./ +RUN npm ci +COPY src ./src +RUN npm run build + +FROM node:24-alpine +ENV NODE_ENV=production +RUN addgroup -S librechat-code && adduser -S librechat-code -G librechat-code +WORKDIR /app +COPY --from=build /app/package.json ./package.json +COPY --from=build /app/dist ./dist +USER librechat-code +ENTRYPOINT ["node", "dist/cli.js"] diff --git a/packages/code/README.md b/packages/code/README.md new file mode 100644 index 00000000..883b2695 --- /dev/null +++ b/packages/code/README.md @@ -0,0 +1,87 @@ +# `@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. + +## Pair + +Hardened deployments use a one-time code instead of copying a long-lived +worker secret onto the VM. After an administrator creates a code, run: + +```bash +librechat-code pair https://code.example.com/v1 '' \ + --worker-id my-vm +``` + +The CLI generates an Ed25519 key locally and writes its paired identity to +`~/.config/librechat/code/my-vm.json` with owner-only permissions. The private +key never leaves the VM. Worker requests carry an exact-request signature, +timestamp, and one-time nonce; the short-lived credential rotates +automatically. + +Then start the worker without a shared secret: + +```bash +LIBRECHAT_CODE_WORKER_ID=my-vm \ +LIBRECHAT_CODE_SANDBOX_ENDPOINT=http://127.0.0.1:2000/api/v2 \ +librechat-code run +``` + +Use `--identity ` while pairing and +`LIBRECHAT_CODE_IDENTITY_FILE=` while running to override the identity +file location. + +## Static compatibility mode + +Non-hardened development deployments may still run with a static token: + +```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 run +``` + +Optional environment variables: + +- `LIBRECHAT_CODE_SANDBOX_PROFILE`: capability label; defaults to `nsjail`. +- `LIBRECHAT_CODE_RUNTIMES`: comma-separated capability labels. +- `LIBRECHAT_CODE_POLICY`: local policy description hashed into the worker's + registration; defaults to `default-deny`. +- `LIBRECHAT_CODE_STATEFUL_WORKSPACE`: defaults to `false`. Set it to `true` + only when the local sandbox supervisor provides a distinct persistent runner + for every runtime session. In that mode the endpoint must contain a + `{runtimeSessionId}` placeholder, for example + `http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2`. The worker URL- + encodes and substitutes the assigned session ID before execution. Hintless + assignments use an ephemeral `assignment-` session so affinity-mode + stateless work never reaches a literal placeholder route. + +A single built-in sandbox runner binds itself to one runtime session and must +not be advertised as stateful. Use the default stateless capability until a +session-routing supervisor is configured. + +Static worker authentication is rejected when Code API hardened mode is +enabled. Expose only the sandbox loopback endpoint to the CLI, and enforce +VM/container egress policy independently of the bridge transport. + +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. + +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. Run it while the normal worker +process is stopped, then restart the normal worker after the command exits. 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.test.ts b/packages/code/src/cli.test.ts new file mode 100644 index 00000000..179bf813 --- /dev/null +++ b/packages/code/src/cli.test.ts @@ -0,0 +1,49 @@ +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/, + ); +}); + +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 new file mode 100644 index 00000000..7edd8a26 --- /dev/null +++ b/packages/code/src/cli.ts @@ -0,0 +1,176 @@ +#!/usr/bin/env node +import { createHash } from 'node:crypto'; + +import { pairBridgeWorker } from './pairing.js'; +import { + defaultBridgeIdentityPath, + loadBridgeIdentity, + saveBridgeIdentity, +} from './storage.js'; +import { BridgeWorker } from './worker.js'; +import { + isValidBridgeWorkerCapabilities, + isValidBridgeWorkerId, +} from './protocol.js'; + +function required(name: string, value = process.env[name]): string { + const normalized = value?.trim(); + if (!normalized) throw new Error(`${name} is required`); + return normalized; +} + +function list(value: string | undefined): string[] { + return ( + value + ?.split(',') + .map((item) => item.trim()) + .filter(Boolean) ?? [] + ); +} + +function option(args: string[], name: string): string | undefined { + const index = args.indexOf(name); + if (index >= 0) return args[index + 1]; + return args.find((value) => value.startsWith(`${name}=`))?.slice(name.length + 1); +} + +async function pair(args: string[]): Promise { + const codeApiUrl = required('instance URL', args[1]); + const code = required('one-time pairing code', args[2]); + const workerId = required( + '--worker-id or LIBRECHAT_CODE_WORKER_ID', + option(args, '--worker-id') ?? process.env.LIBRECHAT_CODE_WORKER_ID, + ); + const identityPath = + option(args, '--identity') ?? + process.env.LIBRECHAT_CODE_IDENTITY_FILE ?? + defaultBridgeIdentityPath(workerId); + const identity = await pairBridgeWorker({ codeApiUrl, workerId, code }); + await saveBridgeIdentity(identityPath, identity); + process.stdout.write( + `Paired worker ${workerId}. Identity saved to ${identityPath}\n`, + ); +} +async function run(runtimeSessionId?: string): Promise { + const configuredWorkerId = process.env.LIBRECHAT_CODE_WORKER_ID?.trim(); + const configuredIdentityPath = process.env.LIBRECHAT_CODE_IDENTITY_FILE?.trim(); + const configuredToken = process.env.LIBRECHAT_CODE_WORKER_TOKEN?.trim(); + const identityPath = + configuredIdentityPath ?? + (configuredWorkerId && !configuredToken + ? defaultBridgeIdentityPath(configuredWorkerId) + : undefined); + const pairedIdentity = identityPath + ? await loadBridgeIdentity(identityPath) + : undefined; + const workerId = required( + 'LIBRECHAT_CODE_WORKER_ID', + configuredWorkerId ?? pairedIdentity?.workerId, + ); + if (!isValidBridgeWorkerId(workerId)) { + throw new Error( + 'LIBRECHAT_CODE_WORKER_ID must match the bridge worker ID format', + ); + } + if (pairedIdentity && pairedIdentity.workerId !== workerId) { + throw new Error( + `Identity belongs to ${pairedIdentity.workerId}, not configured worker ${workerId}`, + ); + } + const codeApiUrl = required( + 'LIBRECHAT_CODE_URL', + process.env.LIBRECHAT_CODE_URL ?? pairedIdentity?.codeApiUrl, + ); + 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 workerIdentity = pairedIdentity + ? { + privateKey: pairedIdentity.privateKey, + credential: pairedIdentity.credential, + expiresAt: pairedIdentity.expiresAt, + } + : undefined; + 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 controller = new AbortController(); + process.once('SIGINT', () => controller.abort()); + process.once('SIGTERM', () => controller.abort()); + const worker = new BridgeWorker({ + codeApiUrl, + token: configuredToken, + identity: workerIdentity, + workerId, + sandboxEndpoint, + capabilities, + onIdentityChange: + pairedIdentity && identityPath + ? async (identity) => { + await saveBridgeIdentity(identityPath, { + ...pairedIdentity, + credential: identity.credential, + expiresAt: identity.expiresAt, + }); + } + : undefined, + onError: (error) => { + const message = + error instanceof Error ? error.message : 'unknown bridge error'; + process.stderr.write(`librechat-code: reconnecting after ${message}\n`); + }, + }); + if (runtimeSessionId !== undefined) { + await worker.refreshCredential(controller.signal); + await worker.register(controller.signal); + await worker.resetWorkspace(runtimeSessionId, controller.signal); + process.stdout.write( + `librechat-code: reset acknowledged for ${runtimeSessionId}\n`, + ); + return; + } + await worker.run(controller.signal); +} + +async function main(): Promise { + const args = process.argv.slice(2); + if (args[0] === 'pair') { + await pair(args); + return; + } + if (args[0] === 'reset-workspace') { + const runtimeSessionId = args[1]?.trim(); + if (!runtimeSessionId) { + throw new Error( + 'Usage: librechat-code reset-workspace ', + ); + } + await run(runtimeSessionId); + return; + } + if (args[0] && args[0] !== 'run') { + throw new Error(`Unknown command: ${args[0]}`); + } + await run(); +} +main().catch((error: Error) => { + process.stderr.write(`librechat-code: ${error.message}\n`); + process.exitCode = 1; +}); diff --git a/packages/code/src/identity.test.ts b/packages/code/src/identity.test.ts new file mode 100644 index 00000000..4e92b91f --- /dev/null +++ b/packages/code/src/identity.test.ts @@ -0,0 +1,35 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { + createBridgeIdentity, + signBridgeRequest, + verifyBridgeRequest, +} from './identity.js'; + +test('worker identity proves possession for the exact HTTP request', () => { + const identity = createBridgeIdentity(); + const request = { + credential: 'short-lived-credential', + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce: 'single-use-request-nonce', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + + const signature = signBridgeRequest(identity.privateKey, request); + + assert.equal( + verifyBridgeRequest(identity.publicKey, request, signature), + true, + ); + assert.equal( + verifyBridgeRequest( + identity.publicKey, + { ...request, body: JSON.stringify({ protocolVersion: 1, waitMs: 0 }) }, + signature, + ), + false, + ); +}); diff --git a/packages/code/src/identity.ts b/packages/code/src/identity.ts new file mode 100644 index 00000000..ab11c865 --- /dev/null +++ b/packages/code/src/identity.ts @@ -0,0 +1,66 @@ +import { + createHash, + generateKeyPairSync, + sign, + verify, +} from 'node:crypto'; + +export interface BridgeIdentity { + publicKey: string; + privateKey: string; +} + +export interface BridgeRequestProofInput { + credential: string; + method: string; + path: string; + timestamp: string; + nonce: string; + body: string; +} + +export function createBridgeIdentity(): BridgeIdentity { + const { publicKey, privateKey } = generateKeyPairSync('ed25519', { + publicKeyEncoding: { type: 'spki', format: 'pem' }, + privateKeyEncoding: { type: 'pkcs8', format: 'pem' }, + }); + return { publicKey, privateKey }; +} + +function canonicalBridgeRequest(input: BridgeRequestProofInput): string { + const bodyDigest = createHash('sha256').update(input.body).digest('hex'); + return [ + input.method.toUpperCase(), + input.path, + input.timestamp, + input.nonce, + bodyDigest, + input.credential, + ].join('\n'); +} + +export function signBridgeRequest( + privateKey: string, + input: BridgeRequestProofInput, +): string { + return sign(null, Buffer.from(canonicalBridgeRequest(input)), privateKey).toString( + 'base64url', + ); +} + +export function verifyBridgeRequest( + publicKey: string, + input: BridgeRequestProofInput, + signature: string, +): boolean { + try { + return verify( + null, + Buffer.from(canonicalBridgeRequest(input)), + publicKey, + Buffer.from(signature, 'base64url'), + ); + } catch { + return false; + } +} diff --git a/packages/code/src/index.ts b/packages/code/src/index.ts new file mode 100644 index 00000000..c65b9f15 --- /dev/null +++ b/packages/code/src/index.ts @@ -0,0 +1,5 @@ +export * from './protocol.js'; +export * from './identity.js'; +export * from './pairing.js'; +export * from './storage.js'; +export * from './worker.js'; diff --git a/packages/code/src/pairing.test.ts b/packages/code/src/pairing.test.ts new file mode 100644 index 00000000..2b564394 --- /dev/null +++ b/packages/code/src/pairing.test.ts @@ -0,0 +1,36 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { pairBridgeWorker } from './pairing.js'; + +test('pairing binds a generated worker key to a single-use code', async () => { + const fetchImpl: typeof fetch = async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { + workerId: string; + code: string; + publicKey: string; + }; + assert.equal(body.workerId, 'vm-1'); + assert.equal(body.code, 'one-time-code'); + assert.match(body.publicKey, /BEGIN PUBLIC KEY/); + return Response.json({ + protocolVersion: 1, + workerId: body.workerId, + credential: 'issued-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }); + }; + + const paired = await pairBridgeWorker({ + codeApiUrl: 'https://code.example/v1/', + workerId: 'vm-1', + code: 'one-time-code', + fetchImpl, + }); + + assert.equal(paired.workerId, 'vm-1'); + assert.equal(paired.codeApiUrl, 'https://code.example/v1'); + assert.equal(paired.credential, 'issued-short-lived-credential-value'); + assert.match(paired.publicKey, /BEGIN PUBLIC KEY/); + assert.match(paired.privateKey, /BEGIN PRIVATE KEY/); +}); diff --git a/packages/code/src/pairing.ts b/packages/code/src/pairing.ts new file mode 100644 index 00000000..ed33fb18 --- /dev/null +++ b/packages/code/src/pairing.ts @@ -0,0 +1,73 @@ +import { createBridgeIdentity } from './identity.js'; +import { + BRIDGE_PROTOCOL_VERSION, + BridgeProtocolError, +} from './protocol.js'; + +import type { BridgeWorkerCredentialResponse } from './protocol.js'; + +export interface PairBridgeWorkerOptions { + codeApiUrl: string; + workerId: string; + code: string; + fetchImpl?: typeof fetch; +} + +export interface PairedBridgeWorkerIdentity + extends BridgeWorkerCredentialResponse { + codeApiUrl: string; + publicKey: string; + privateKey: string; +} + +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 async function pairBridgeWorker( + options: PairBridgeWorkerOptions, +): Promise { + const codeApiUrl = normalizedBaseUrl(options.codeApiUrl); + const identity = createBridgeIdentity(); + const response = await (options.fetchImpl ?? fetch)( + `${codeApiUrl}/bridge/pairings/redeem`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: options.workerId, + code: options.code, + publicKey: identity.publicKey, + }), + }, + ); + const payload = (await response.json()) as object; + if (!response.ok) { + throw new BridgeProtocolError( + errorMessage(payload) ?? `Bridge pairing failed with HTTP ${response.status}`, + response.status, + ); + } + const credential = payload as BridgeWorkerCredentialResponse; + if ( + credential.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + credential.workerId !== options.workerId || + typeof credential.credential !== 'string' || + credential.credential.length < 32 || + !Number.isFinite(Date.parse(credential.expiresAt)) + ) { + throw new BridgeProtocolError('Code API returned an invalid worker credential'); + } + return { + ...credential, + codeApiUrl, + publicKey: identity.publicKey, + privateKey: identity.privateKey, + }; +} diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts new file mode 100644 index 00000000..61c50922 --- /dev/null +++ b/packages/code/src/protocol.test.ts @@ -0,0 +1,56 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import { + bridgeWorkerPath, + isValidBridgeWorkerCapabilities, + isValidBridgeWorkerId, +} from './protocol.js'; + +test('bridgeWorkerPath encodes worker-controlled path segments', () => { + assert.equal( + bridgeWorkerPath('vm/example worker'), + '/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); +}); + +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 new file mode 100644 index 00000000..b2027e1c --- /dev/null +++ b/packages/code/src/protocol.ts @@ -0,0 +1,143 @@ +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; + +export interface BridgeWorkerCapabilities { + statefulWorkspace: boolean; + sandboxProfile: string; + runtimes: string[]; + policyDigest?: string; +} + +export interface BridgeWorkerRegistration { + protocolVersion: BridgeProtocolVersion; + workerId: string; + incarnationId: string; + capabilities: BridgeWorkerCapabilities; +} + +export interface BridgeWorkerRegistrationResponse { + protocolVersion: BridgeProtocolVersion; + workerId: string; + incarnationId: string; + registeredAt: string; + leaseTtlMs: number; +} + +export interface BridgePairingRedemption { + protocolVersion: BridgeProtocolVersion; + workerId: string; + code: string; + publicKey: string; +} + +export interface BridgeWorkerCredentialResponse { + protocolVersion: BridgeProtocolVersion; + workerId: string; + credential: string; + expiresAt: string; +} + +export interface BridgeSandboxRequest { + body: TBody; + headers: Record; +} + +export interface BridgeAssignment { + protocolVersion: BridgeProtocolVersion; + assignmentId: string; + workerId: string; + incarnationId: string; + generation: number; + leaseToken: string; + expiresAt: string; + /** Server-calculated execution budget at lease time; avoids VM clock skew. */ + remainingMs?: number; + runtimeSessionId?: string; + request: BridgeSandboxRequest; +} + +export interface BridgeLeaseResponse { + protocolVersion: BridgeProtocolVersion; + /** Time spent handling the lease request on Code API, excluding transit. */ + serverElapsedMs?: number; + assignment?: BridgeAssignment; +} + +export interface BridgeFulfilledSettlement { + protocolVersion: BridgeProtocolVersion; + generation: number; + leaseToken: string; + incarnationId: string; + status: 'fulfilled'; + result: TResult; +} + +export interface BridgeRejectedSettlement { + protocolVersion: BridgeProtocolVersion; + generation: number; + leaseToken: string; + incarnationId: string; + status: 'rejected'; + error: string; +} + +export type BridgeSettlement = + BridgeFulfilledSettlement | BridgeRejectedSettlement; + +export interface BridgeSettlementResponse { + protocolVersion: BridgeProtocolVersion; + accepted: true; +} + +export interface BridgeCancellationResponse { + protocolVersion: BridgeProtocolVersion; + cancelled: boolean; +} + +export class BridgeProtocolError extends Error { + constructor( + message: string, + public readonly status?: number, + public readonly code?: string, + ) { + super(message); + this.name = 'BridgeProtocolError'; + } +} + +export function bridgeWorkerPath(workerId: string): string { + return `/bridge/workers/${encodeURIComponent(workerId)}`; +} + +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/packages/code/src/storage.test.ts b/packages/code/src/storage.test.ts new file mode 100644 index 00000000..ddd782b0 --- /dev/null +++ b/packages/code/src/storage.test.ts @@ -0,0 +1,41 @@ +import assert from 'node:assert/strict'; +import { mkdtemp, rm, stat } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; + +import { + defaultBridgeIdentityPath, + loadBridgeIdentity, + saveBridgeIdentity, +} from './storage.js'; + +test('default identity paths do not collide after worker ID sanitization', () => { + assert.notEqual( + defaultBridgeIdentityPath('vm:a'), + defaultBridgeIdentityPath('vm_a'), + ); +}); + +test('paired identity is persisted atomically with owner-only permissions', async () => { + const directory = await mkdtemp(join(tmpdir(), 'librechat-code-')); + const path = join(directory, 'identity.json'); + const identity = { + protocolVersion: 1 as const, + workerId: 'vm-1', + codeApiUrl: 'https://code.example/v1', + credential: 'issued-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + publicKey: 'public-key', + privateKey: 'private-key', + }; + + try { + await saveBridgeIdentity(path, identity); + + assert.deepEqual(await loadBridgeIdentity(path), identity); + assert.equal((await stat(path)).mode & 0o777, 0o600); + } finally { + await rm(directory, { recursive: true, force: true }); + } +}); diff --git a/packages/code/src/storage.ts b/packages/code/src/storage.ts new file mode 100644 index 00000000..a26c3eb0 --- /dev/null +++ b/packages/code/src/storage.ts @@ -0,0 +1,66 @@ +import { createHash, randomBytes } from 'node:crypto'; +import { chmod, mkdir, open, readFile, rename, rm } from 'node:fs/promises'; +import { homedir } from 'node:os'; +import { dirname, join } from 'node:path'; + +import { BRIDGE_PROTOCOL_VERSION, BridgeProtocolError } from './protocol.js'; + +import type { PairedBridgeWorkerIdentity } from './pairing.js'; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function isPairedIdentity(value: unknown): value is PairedBridgeWorkerIdentity { + if (!isRecord(value)) return false; + return ( + value.protocolVersion === BRIDGE_PROTOCOL_VERSION && + typeof value.workerId === 'string' && + typeof value.codeApiUrl === 'string' && + typeof value.credential === 'string' && + typeof value.expiresAt === 'string' && + Number.isFinite(Date.parse(value.expiresAt)) && + typeof value.publicKey === 'string' && + typeof value.privateKey === 'string' + ); +} + +export function defaultBridgeIdentityPath(workerId: string): string { + const readableName = workerId.replace(/[^A-Za-z0-9._-]/g, '_'); + const fileName = readableName === workerId + ? readableName + : `${readableName}-${createHash('sha256').update(workerId).digest('hex').slice(0, 16)}`; + return join(homedir(), '.config', 'librechat', 'code', `${fileName}.json`); +} + +export async function saveBridgeIdentity( + path: string, + identity: PairedBridgeWorkerIdentity, +): Promise { + await mkdir(dirname(path), { recursive: true, mode: 0o700 }); + const temporaryPath = `${path}.${randomBytes(8).toString('hex')}.tmp`; + try { + const file = await open(temporaryPath, 'wx', 0o600); + try { + await file.writeFile(`${JSON.stringify(identity, null, 2)}\n`, 'utf8'); + await file.sync(); + } finally { + await file.close(); + } + await rename(temporaryPath, path); + await chmod(path, 0o600); + } catch (error) { + await rm(temporaryPath, { force: true }); + throw error; + } +} + +export async function loadBridgeIdentity( + path: string, +): Promise { + const identity = JSON.parse(await readFile(path, 'utf8')) as unknown; + if (!isPairedIdentity(identity)) { + throw new BridgeProtocolError(`Invalid bridge identity file: ${path}`); + } + return identity; +} diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts new file mode 100644 index 00000000..f4ec7d0b --- /dev/null +++ b/packages/code/src/worker.test.ts @@ -0,0 +1,2591 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import { createBridgeIdentity, verifyBridgeRequest } from './identity.js'; +import { + BridgeWorker, + BridgeWorkspaceQuarantinedError, + reconnectDelayMs, +} from './worker.js'; + +import type { BridgeAssignment } from './protocol.js'; + +const incarnationId = 'incarnation-00000001'; + +test('worker forwards a fenced assignment to the sandbox and settles the result', async () => { + const requests: Array<{ url: string; init?: RequestInit }> = []; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + requests.push({ url, init }); + if (url.endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1/', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2/', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + const assignment: BridgeAssignment = { + protocolVersion: 1, + assignmentId: 'assignment-1', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 3, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { + body: { language: 'bash' }, + headers: { 'X-Execution-Manifest': 'signed' }, + }, + }; + + await worker.executeAndSettle(assignment); + + assert.equal(requests.length, 2); + assert.equal( + requests[0].url, + 'http://127.0.0.1:2000/sessions/rt-user-1/api/v2/execute', + ); + assert.equal( + (requests[0].init?.headers as Record)[ + 'X-Runtime-Session-Id' + ], + 'rt-user-1', + ); + assert.match(requests[1].url, /assignments\/assignment-1\/settle$/); + assert.deepEqual(JSON.parse(String(requests[1].init?.body)), { + protocolVersion: 1, + generation: 3, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + incarnationId: 'incarnation-00000001', + status: 'fulfilled', + result: { session_id: 'run-1', files: [] }, + }); +}); + +test('worker 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 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, + 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, serverElapsedMs: 0, assignment }), + { 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' }, + }); + } + 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) => { + 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: 100, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + if (url.endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 20)); + 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 new Promise((resolve) => setTimeout(resolve, 45)); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-heartbeat', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.ok(registrations >= 2); +}); + +test('worker 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; + 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 = ''; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + executeUrl = url; + runtimeSessionHeader = (init?.headers as Record)[ + 'X-Runtime-Session-Id' + ]; + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'hintless-assignment', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal( + executeUrl, + 'http://127.0.0.1:2000/sessions/assignment-hintless-assignment/api/v2/execute', + ); + assert.equal(runtimeSessionHeader, 'assignment-hintless-assignment'); +}); + +test('worker quarantines a fulfilled stateful settlement rejected by Code API', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response(JSON.stringify({ error: 'assignment was fenced' }), { + status: 409, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'fenced-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); +}); + +test('worker surfaces a definite stateless settlement rejection directly', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + return new Response(JSON.stringify({ error: 'assignment was fenced' }), { + status: 409, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'fenced-stateless-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }), + (error: unknown) => + error instanceof Error && + error.name === 'BridgeProtocolError' && + error.message === 'assignment was fenced', + ); +}); + +test('worker 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) => { + const url = String(input); + if (url.endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + settlementAttempts += 1; + if (settlementAttempts === 1) throw new TypeError('connection reset'); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'retry-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(settlementAttempts, 2); +}); + +test('worker quarantines stateful reuse after settlement stays ambiguous', async () => { + const fetchImpl: typeof fetch = async (input) => { + if (String(input).endsWith('/execute')) { + return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + } + throw new TypeError('connection reset'); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'ambiguous-settlement', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 50).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); +}); + +test('worker 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, + rejectionAckGraceMs: 0, + }); + + 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 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, + 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.register(); + 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); + assert.ok(registrations > 1); +}); + +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', + 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 === 2) { + 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' }, + }); + } + 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' } }, + ); + }, + }); + + 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 >= 3); + assert.equal(rejectedSettlement, true); + assert.equal(settlementAttempts, 2); +}); + +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 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) => { + if (String(input).endsWith('/execute')) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + settlementAttempted = true; + return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'aborted-execution', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 30).toISOString(), + runtimeSessionId: 'rt-user-1', + request: { body: { language: 'bash' }, headers: {} }, + }), + BridgeWorkspaceQuarantinedError, + ); + assert.equal(settlementAttempted, false); +}); + +test('worker surfaces quarantine when shutdown aborts stateful execution', async () => { + const controller = new AbortController(); + let executeStarted = false; + const assignment: BridgeAssignment = { + protocolVersion: 1, + assignmentId: 'shutdown-execution', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + remainingMs: 5_000, + 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, serverElapsedMs: 0, 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); +}); + +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 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', + 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 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 (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({ + 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 rejects a lease whose acknowledgement exhausts its budget', async () => { + const originalNow = Date.now; + let now = 100_000; + let abandonedSettlement: Record | undefined; + let registrations = 0; + let settlementAttempts = 0; + 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, 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( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + 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' } }, + ); + } + 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/); + assert.equal(abandonedSettlement?.status, 'rejected'); + assert.ok(registrations > 0); + assert.equal(settlementAttempts, 2); + } finally { + Date.now = originalNow; + } +}); + +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({ + 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({ + 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); +}); + +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 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({ + 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); +}); + + +test('worker continues after an expired assignment settlement conflict', async () => { + const controller = new AbortController(); + let registrations = 0; + let leases = 0; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + reconnectDelayMs: 0, + reconnectMaxDelayMs: 0, + fetchImpl: async (input) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + registrations += 1; + if (registrations === 2) controller.abort(); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + } + if (url.endsWith('/lease')) { + leases += 1; + return Response.json({ + protocolVersion: 1, + assignment: leases === 1 + ? { + protocolVersion: 1, + assignmentId: 'assignment-expired', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + } + : undefined, + }); + } + if (url.endsWith('/execute')) { + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/settle')) { + return Response.json( + { error: 'Bridge assignment has expired', code: 'ASSIGNMENT_EXPIRED' }, + { status: 409 }, + ); + } + return Response.json({ cancelled: false }); + }, + }); + + await worker.run(controller.signal); + + assert.equal(registrations, 2); +}); + +test('paired worker proves possession on bridge requests', async () => { + const key = createBridgeIdentity(); + let bridgeRequest: { url: string; init?: RequestInit } | undefined; + const fetchImpl: typeof fetch = async (input, init) => { + bridgeRequest = { url: String(input), init }; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'issued-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.register(); + + assert.ok(bridgeRequest); + const headers = bridgeRequest.init?.headers as Record; + const body = String(bridgeRequest.init?.body); + assert.equal( + verifyBridgeRequest( + key.publicKey, + { + credential: 'issued-short-lived-credential-value', + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: headers['X-LibreChat-Code-Timestamp'], + nonce: headers['X-LibreChat-Code-Nonce'], + body, + }, + headers['X-LibreChat-Code-Signature'], + ), + true, + ); +}); + +test('paired worker rotates an expiring credential before registration', async () => { + const key = createBridgeIdentity(); + const requests: Array<{ url: string; init?: RequestInit }> = []; + let persistedCredential = ''; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + requests.push({ url, init }); + if (url.endsWith('/credentials/refresh')) { + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'rotated-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }); + } + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'original-short-lived-credential-value', + expiresAt: new Date(Date.now() + 30_000).toISOString(), + }, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + onIdentityChange: (identity) => { + persistedCredential = identity.credential; + }, + }); + + await worker.refreshCredential(); + await worker.register(); + + assert.equal(persistedCredential, 'rotated-short-lived-credential-value'); + assert.equal( + (requests[1].init?.headers as Record).Authorization, + 'Bridge rotated-short-lived-credential-value', + ); +}); + +test('paired worker retries persistence before adopting a rotated credential', async () => { + const key = createBridgeIdentity(); + const identity = { + privateKey: key.privateKey, + credential: 'original-short-lived-credential-value', + expiresAt: new Date(Date.now() + 30_000).toISOString(), + }; + let persistenceAttempts = 0; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async () => + Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'rotated-short-lived-credential-value', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }), + onIdentityChange: () => { + persistenceAttempts += 1; + if (persistenceAttempts === 1) throw new Error('disk unavailable'); + }, + }); + + await assert.rejects(worker.refreshCredential(), /disk unavailable/); + assert.equal(identity.credential, 'original-short-lived-credential-value'); + await worker.refreshCredential(); + assert.equal(identity.credential, 'rotated-short-lived-credential-value'); + assert.equal(persistenceAttempts, 2); +}); + +test('paired worker refreshes before an assignment that outlives its credential', async () => { + const key = createBridgeIdentity(); + const requests: string[] = []; + const fetchImpl: typeof fetch = async (input) => { + const url = String(input); + requests.push(url); + if (url.endsWith('/credentials/refresh')) { + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'assignment-safe-rotated-credential-value', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }); + } + if (url.endsWith('/execute')) { + return Response.json({ session_id: 'run-long', files: [] }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-too-short-for-assignment', + expiresAt: new Date(Date.now() + 30_000).toISOString(), + }, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-long', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'assignment-long-lease-token-value', + expiresAt: new Date(Date.now() + 5 * 60_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.match(requests[0], /credentials\/refresh$/); + assert.equal(requests[1], 'http://127.0.0.1:2000/api/v2/execute'); +}); + +test('paired worker rotates credentials throughout a long assignment', async () => { + const key = createBridgeIdentity(); + let refreshCount = 0; + const identity = { + privateKey: key.privateKey, + credential: 'credential-before-long-running-assignment', + expiresAt: new Date(Date.now() + 5).toISOString(), + }; + const fetchImpl: typeof fetch = async (input) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + refreshCount += 1; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: `rotated-long-assignment-credential-${refreshCount}`, + expiresAt: new Date(Date.now() + 30).toISOString(), + }); + } + if (url.endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 55)); + return Response.json({ session_id: 'run-long-rotation', files: [] }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-credential-maintenance', + workerId: 'vm-1', + incarnationId, + generation: 5, + leaseToken: 'assignment-credential-maintenance-token', + expiresAt: new Date(Date.now() + 500).toISOString(), + remainingMs: 500, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.ok(refreshCount >= 2); + assert.match(identity.credential, /^rotated-long-assignment-credential-/); +}); + +test('paired worker cancels a stalled credential refresh after execution', async () => { + const key = createBridgeIdentity(); + let refreshStarted!: () => void; + const started = new Promise((resolve) => { + refreshStarted = resolve; + }); + let refreshAborted = false; + const identity = { + privateKey: key.privateKey, + credential: 'credential-before-stalled-refresh', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + credentialRefreshWindowMs: 100, + credentialRefreshTransportTimeoutMs: 10_000, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + refreshStarted(); + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => { + refreshAborted = true; + reject(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + }); + } + if (url.endsWith('/execute')) { + await started; + return Response.json({ + session_id: 'run-stalled-refresh', + files: [], + }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + queueMicrotask(() => { + identity.expiresAt = new Date(Date.now() + 50).toISOString(); + }); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-stalled-refresh', + workerId: 'vm-1', + incarnationId, + generation: 6, + leaseToken: 'assignment-stalled-refresh-token', + expiresAt: new Date(Date.now() + 2_000).toISOString(), + remainingMs: 2_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(refreshAborted, true); +}); + +test('paired worker refreshes conservatively before server clock calibration', async () => { + const key = createBridgeIdentity(); + let refreshCount = 0; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-idle-clock-skew-refresh', + expiresAt: new Date(Date.now() + 65_000).toISOString(), + }, + credentialRefreshWindowMs: 10_000, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async () => { + refreshCount += 1; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-idle-clock-skew-refresh', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }); + }, + }); + + await worker.refreshCredential(); + + assert.equal(refreshCount, 1); +}); + +test('paired worker charges initial credential refresh against the assignment deadline', async () => { + const key = createBridgeIdentity(); + let sandboxStarted = false; + let rejected = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-deadline-refresh', + expiresAt: new Date(Date.now() + 5).toISOString(), + }, + credentialRefreshWindowMs: 10, + credentialRefreshTransportTimeoutMs: 10_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + if (url.endsWith('/execute')) { + sandboxStarted = true; + } + if (url.endsWith('/settle')) { + rejected = + JSON.parse(String(init?.body)).status === 'rejected'; + } + return Response.json({ + protocolVersion: 1, + accepted: true, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-deadline-refresh', + workerId: 'vm-1', + incarnationId, + generation: 7, + leaseToken: 'assignment-deadline-refresh-token', + expiresAt: new Date(Date.now() + 30).toISOString(), + remainingMs: 30, + runtimeSessionId: 'rt-deadline-refresh', + request: { body: { language: 'bash' }, headers: {} }, + }), + { name: 'AbortError' }, + ); + + assert.equal(sandboxStarted, false); + assert.equal(rejected, true); +}); + +test('paired worker rechecks the deadline after request serialization', async () => { + const key = createBridgeIdentity(); + let sandboxStarted = false; + let rejected = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-valid-during-serialization', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + sandboxStarted = true; + } + if (url.endsWith('/settle')) { + rejected = + JSON.parse(String(init?.body)).status === 'rejected'; + } + return Response.json({ + protocolVersion: 1, + accepted: true, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }, + }); + const slowBody = { + get language(): string { + const blockedUntilMs = Date.now() + 25; + while (Date.now() < blockedUntilMs) { + // Deliberately consume the remaining synchronous request budget. + } + return 'bash'; + }, + }; + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-serialization-deadline', + workerId: 'vm-1', + incarnationId, + generation: 8, + leaseToken: 'assignment-serialization-deadline-token', + expiresAt: new Date(Date.now() + 10).toISOString(), + remainingMs: 10, + runtimeSessionId: 'rt-serialization-deadline', + request: { body: slowBody, headers: {} }, + }); + + assert.equal(sandboxStarted, false); + assert.equal(rejected, true); +}); + +test('paired worker keeps endpoint validation failures known-clean', async () => { + const key = createBridgeIdentity(); + let sandboxStarted = false; + let rejected = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-for-invalid-endpoint', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) sandboxStarted = true; + if (url.endsWith('/settle')) { + rejected = + JSON.parse(String(init?.body)).status === 'rejected'; + } + return Response.json({ + protocolVersion: 1, + accepted: true, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-invalid-endpoint', + workerId: 'vm-1', + incarnationId, + generation: 9, + leaseToken: 'assignment-invalid-endpoint-token', + expiresAt: new Date(Date.now() + 500).toISOString(), + remainingMs: 500, + runtimeSessionId: 'rt-invalid-endpoint', + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(sandboxStarted, false); + assert.equal(rejected, true); +}); + +test('paired worker rechecks shutdown after persisting a refreshed identity', async () => { + const key = createBridgeIdentity(); + const controller = new AbortController(); + let sandboxStarted = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-shutdown-refresh', + expiresAt: new Date(Date.now() + 5).toISOString(), + }, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input) => { + if (String(input).endsWith('/execute')) sandboxStarted = true; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-persisted-during-shutdown', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }); + }, + onIdentityChange: () => { + controller.abort(); + }, + }); + + await assert.rejects( + worker.executeAndSettle( + { + protocolVersion: 1, + assignmentId: 'assignment-shutdown-refresh', + workerId: 'vm-1', + incarnationId, + generation: 10, + leaseToken: 'assignment-shutdown-refresh-token', + expiresAt: new Date(Date.now() + 500).toISOString(), + remainingMs: 500, + runtimeSessionId: 'rt-shutdown-refresh', + request: { body: { language: 'bash' }, headers: {} }, + }, + controller.signal, + ), + { name: 'AbortError' }, + ); + + assert.equal(sandboxStarted, false); +}); + +test('paired worker retries transient refresh failures before credential expiry', async () => { + const key = createBridgeIdentity(); + let refreshCount = 0; + const identity = { + privateKey: key.privateKey, + credential: 'credential-before-transient-refresh', + expiresAt: new Date(Date.now() + 40).toISOString(), + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + credentialRefreshWindowMs: 15, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + refreshCount += 1; + if (refreshCount === 1) { + return Response.json( + { error: 'temporarily unavailable' }, + { status: 503 }, + ); + } + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-transient-refresh', + expiresAt: new Date(Date.now() + 500).toISOString(), + }); + } + if (url.endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 70)); + return Response.json({ + session_id: 'run-transient-refresh', + files: [], + }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-transient-refresh', + workerId: 'vm-1', + incarnationId, + generation: 7, + leaseToken: 'assignment-transient-refresh-token', + expiresAt: new Date(Date.now() + 500).toISOString(), + remainingMs: 500, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(refreshCount, 2); + assert.equal(identity.credential, 'credential-after-transient-refresh'); +}); + +test('paired worker preserves its refresh margin when the server clock is ahead', async () => { + const key = createBridgeIdentity(); + const serverClockOffsetMs = 55; + let refreshCount = 0; + const identity = { + privateKey: key.privateKey, + credential: 'credential-before-clock-skew-refresh', + expiresAt: new Date(Date.now() + serverClockOffsetMs + 20).toISOString(), + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + refreshCount += 1; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-clock-skew-refresh', + expiresAt: new Date( + Date.now() + serverClockOffsetMs + 500, + ).toISOString(), + }); + } + if (url.endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 35)); + return Response.json({ + session_id: 'run-clock-skew-refresh', + files: [], + }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + const remainingMs = 500; + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-clock-skew-refresh', + workerId: 'vm-1', + incarnationId, + generation: 8, + leaseToken: 'assignment-clock-skew-refresh-token', + expiresAt: new Date( + Date.now() + serverClockOffsetMs + remainingMs, + ).toISOString(), + remainingMs, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(refreshCount, 1); + assert.equal(identity.credential, 'credential-after-clock-skew-refresh'); +}); + +test('worker shutdown interrupts reconnect backoff', async () => { + const controller = new AbortController(); + let failed!: () => void; + const failure = new Promise((resolve) => { + failed = resolve; + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async () => { + throw new Error('offline'); + }, + reconnectDelayMs: 30_000, + reconnectMaxDelayMs: 30_000, + onError: () => failed(), + }); + + const run = worker.run(controller.signal); + await failure; + controller.abort(); + await run; +}); + +test('sandbox completion does not cancel an in-flight credential rotation', async () => { + const key = createBridgeIdentity(); + const identity = { + privateKey: key.privateKey, + credential: 'credential-before-in-flight-rotation', + expiresAt: new Date(Date.now() + 40).toISOString(), + }; + let refreshStarted!: () => void; + const refreshStartedPromise = new Promise((resolve) => { + refreshStarted = resolve; + }); + let refreshCount = 0; + let settleAuthorization = ''; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + refreshCount += 1; + if (refreshCount > 1) { + return Response.json({ error: 'stale credential' }, { status: 401 }); + } + refreshStarted(); + await new Promise((resolve, reject) => { + const timer = setTimeout(resolve, 30); + init?.signal?.addEventListener( + 'abort', + () => { + clearTimeout(timer); + reject(new DOMException('Aborted', 'AbortError')); + }, + { once: true }, + ); + }); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-in-flight-rotation', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }); + } + if (url.endsWith('/execute')) { + await refreshStartedPromise; + return Response.json({ session_id: 'run-rotation-race', files: [] }); + } + settleAuthorization = ( + init?.headers as Record + ).Authorization; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + credentialRefreshWindowMs: 30, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-rotation-race', + workerId: 'vm-1', + incarnationId, + generation: 5, + leaseToken: 'assignment-rotation-race-lease-token', + expiresAt: new Date(Date.now() + 600_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(refreshCount, 1); + assert.equal(identity.credential, 'credential-after-in-flight-rotation'); + assert.equal( + settleAuthorization, + 'Bridge credential-after-in-flight-rotation', + ); +}); + +test('reconnect delay uses bounded exponential jitter', () => { + assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 0), 500); + assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 1), 1_000); + assert.equal(reconnectDelayMs(10, 1_000, 30_000, () => 1), 30_000); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts new file mode 100644 index 00000000..ebbd486e --- /dev/null +++ b/packages/code/src/worker.ts @@ -0,0 +1,1044 @@ +import { randomBytes } from 'node:crypto'; + +import { + BRIDGE_PROTOCOL_VERSION, + BridgeProtocolError, + bridgeWorkerPath, +} from './protocol.js'; +import { signBridgeRequest } from './identity.js'; + +import type { + BridgeAssignment, + BridgeLeaseResponse, + BridgeSettlement, + BridgeSettlementResponse, + BridgeWorkerCapabilities, + BridgeWorkerCredentialResponse, + BridgeWorkerRegistrationResponse, +} from './protocol.js'; + +export interface BridgeWorkerOptions { + codeApiUrl: string; + token?: string; + identity?: BridgeWorkerIdentity; + workerId: string; + sandboxEndpoint: string; + capabilities: BridgeWorkerCapabilities; + leaseWaitMs?: number; + leaseTransportGraceMs?: number; + registrationTransportTimeoutMs?: number; + leaseAckTransportTimeoutMs?: number; + resetTransportTimeoutMs?: number; + cancellationPollIntervalMs?: number; + cancellationTransportTimeoutMs?: number; + rejectionAckGraceMs?: number; + reconnectDelayMs?: number; + reconnectMaxDelayMs?: number; + reconnectRandom?: () => number; + credentialRefreshWindowMs?: number; + credentialRefreshTransportTimeoutMs?: number; + fetchImpl?: typeof fetch; + onError?: (error: unknown) => void; + onIdentityChange?: (identity: BridgeWorkerIdentity) => void | Promise; + incarnationId?: string; +} + +export interface BridgeWorkerIdentity { + privateKey: string; + credential: string; + expiresAt: string; +} + +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_RECONNECT_MAX_DELAY_MS = 30_000; +const CREDENTIAL_REFRESH_WINDOW_MS = 60_000; +const MAX_PROOF_CLOCK_SKEW_MS = 60_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; +const REGISTRATION_RETRY_DELAY_MS = 100; +const CREDENTIAL_REFRESH_RETRY_DELAY_MS = 100; +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}'; + +export function reconnectDelayMs( + attempt: number, + baseDelayMs = DEFAULT_RECONNECT_DELAY_MS, + maxDelayMs = DEFAULT_RECONNECT_MAX_DELAY_MS, + random: () => number = Math.random, +): number { + const cap = Math.min(maxDelayMs, baseDelayMs * 2 ** Math.max(0, attempt)); + return Math.floor(cap * (0.5 + Math.min(1, Math.max(0, random())) * 0.5)); +} + +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; +} + +async function abortableDelay(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted === true) return; + await new Promise((resolve) => { + const onAbort = (): void => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} + +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, + public readonly cause?: unknown, + ) { + super(message); + this.name = 'BridgeWorkspaceQuarantinedError'; + } +} + +export class BridgeWorker { + private readonly fetchImpl: typeof fetch; + private readonly codeApiUrl: string; + private readonly sandboxEndpoint: string; + private readonly incarnationId: string; + private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; + private lastRegisteredAtMs = 0; + private serverClockOffsetMs = MAX_PROOF_CLOCK_SKEW_MS; + + constructor(private readonly options: BridgeWorkerOptions) { + if (!options.token && !options.identity) { + throw new BridgeProtocolError( + 'Bridge worker requires a static token or paired identity', + ); + } + this.fetchImpl = options.fetchImpl ?? fetch; + this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); + this.sandboxEndpoint = normalizedBaseUrl(options.sandboxEndpoint); + this.incarnationId = + options.incarnationId ?? randomBytes(18).toString('base64url'); + } + + async register( + signal?: AbortSignal, + ): Promise { + const 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); + const registrationStartedAtMs = Date.now(); + 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', + ); + } + const registeredAtMs = Date.parse(registration.registeredAt); + if (Number.isFinite(registeredAtMs)) { + this.serverClockOffsetMs = registeredAtMs - registrationStartedAtMs; + } + this.registrationTtlMs = registration.leaseTtlMs; + this.lastRegisteredAtMs = registrationStartedAtMs; + return registration; + } + + async resetWorkspace( + runtimeSessionId: string, + signal?: AbortSignal, + ): Promise { + if (runtimeSessionId.trim().length === 0) { + throw new BridgeProtocolError('Runtime session ID is required'); + } + await this.timedRequest( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + runtimeSessionId, + confirmDiscarded: true, + }, + Math.max( + 1, + this.options.resetTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + ), + signal, + ); + } + + async lease(signal?: AbortSignal): Promise { + 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; + const requestStartedAtMs = Date.now(); + 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 + ) { + throw new BridgeProtocolError( + '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', + ); + } + 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), + ); + const adjustedAssignment = { + ...response.assignment, + remainingMs: Math.max( + 0, + (response.assignment.remainingMs ?? 0) - transportElapsedMs, + ), + }; + const acknowledgementStartedAtMs = Date.now(); + 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) { + await this.rejectUnexecutedAssignment( + adjustedAssignment, + 'Bridge assignment expired during lease acknowledgement', + ); + throw new BridgeProtocolError( + 'Bridge assignment expired during lease acknowledgement', + ); + } + return { + ...adjustedAssignment, + remainingMs, + }; + } + + async run(signal?: AbortSignal): Promise { + let reconnectAttempt = 0; + while (!signal?.aborted) { + try { + await this.refreshCredential(signal); + await this.register(signal); + const assignment = await this.lease(signal); + reconnectAttempt = 0; + if (!assignment) continue; + await this.executeAndSettle(assignment, signal); + } catch (error) { + if (error instanceof BridgeWorkspaceQuarantinedError) { + throw error; + } + if (signal?.aborted) return; + if ( + error instanceof BridgeProtocolError && + (error.status === 401 || + error.status === 403 || + error.code === 'WORKER_FENCED' || + error.code === 'WORKER_QUARANTINED') + ) { + throw error; + } + this.options.onError?.(error); + const delay = reconnectDelayMs( + reconnectAttempt, + this.options.reconnectDelayMs, + this.options.reconnectMaxDelayMs, + this.options.reconnectRandom, + ); + reconnectAttempt += 1; + await abortableDelay(delay, signal); + } + } + } + + async refreshCredential( + signal?: AbortSignal, + validThroughMs = + Date.now() + + this.serverClockOffsetMs + + (this.options.credentialRefreshWindowMs ?? CREDENTIAL_REFRESH_WINDOW_MS), + transportTimeoutMs = Number.POSITIVE_INFINITY, + ): Promise { + const identity = this.options.identity; + if (identity == null) return; + if (Date.parse(identity.expiresAt) > validThroughMs) { + return; + } + const credential = await this.timedRequest( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + + '/credentials/refresh', + { protocolVersion: BRIDGE_PROTOCOL_VERSION }, + Math.max( + 1, + Math.min( + transportTimeoutMs, + this.options.credentialRefreshTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + ), + ), + signal, + ); + if ( + credential.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + credential.workerId !== this.options.workerId || + typeof credential.credential !== 'string' || + credential.credential.length < 32 || + !Number.isFinite(Date.parse(credential.expiresAt)) || + Date.parse(credential.expiresAt) <= validThroughMs + ) { + throw new BridgeProtocolError( + 'Code API returned an invalid rotated worker credential', + ); + } + const rotatedIdentity: BridgeWorkerIdentity = { + ...identity, + credential: credential.credential, + expiresAt: credential.expiresAt, + }; + await this.options.onIdentityChange?.(rotatedIdentity); + identity.credential = rotatedIdentity.credential; + identity.expiresAt = rotatedIdentity.expiresAt; + } + + private async maintainCredential( + assignment: BridgeAssignment, + stopSignal: AbortSignal, + serverClockOffsetMs: number, + requestSignal?: AbortSignal, + ): Promise { + const identity = this.options.identity; + if (identity == null) return; + const refreshWindowMs = + this.options.credentialRefreshWindowMs ?? CREDENTIAL_REFRESH_WINDOW_MS; + const assignmentDeadlineMs = + Date.parse(assignment.expiresAt) - serverClockOffsetMs; + while (!stopSignal.aborted && Date.now() < assignmentDeadlineMs) { + const refreshAtMs = + Date.parse(identity.expiresAt) - serverClockOffsetMs - refreshWindowMs; + const waitMs = Math.max( + 0, + Math.min(refreshAtMs - Date.now(), assignmentDeadlineMs - Date.now()), + ); + await abortableDelay(waitMs, stopSignal); + if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return; + try { + await this.refreshCredential( + requestSignal, + Date.now() + serverClockOffsetMs + refreshWindowMs, + ); + } catch (error) { + if (stopSignal.aborted) return; + const terminal = + error instanceof BridgeProtocolError && + (error.status === 401 || error.status === 403); + const credentialRemainingMs = + Date.parse(identity.expiresAt) - + (Date.now() + serverClockOffsetMs); + if (terminal || credentialRemainingMs <= 0) throw error; + await abortableDelay( + Math.min( + CREDENTIAL_REFRESH_RETRY_DELAY_MS, + Math.max(1, Math.floor(credentialRemainingMs / 2)), + ), + stopSignal, + ); + } + } + } + + async executeAndSettle( + assignment: BridgeAssignment, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted === true) { + throw signal.reason instanceof Error + ? signal.reason + : new DOMException('aborted', 'AbortError'); + } + const serverClockOffsetMs = + Number.isSafeInteger(assignment.remainingMs) && + (assignment.remainingMs ?? -1) >= 0 + ? Date.parse(assignment.expiresAt) - + (Date.now() + (assignment.remainingMs ?? 0)) + : 0; + this.serverClockOffsetMs = serverClockOffsetMs; + const localDeadlineAtMs = + Date.now() + this.assignmentRemainingMs(assignment); + try { + await this.refreshCredential( + signal, + Date.now() + + serverClockOffsetMs + + (this.options.credentialRefreshWindowMs ?? + CREDENTIAL_REFRESH_WINDOW_MS), + Math.max(1, localDeadlineAtMs - Date.now()), + ); + } catch (error) { + if (!signal?.aborted) { + await this.rejectUnexecutedAssignment( + assignment, + 'Bridge credential refresh failed before sandbox execution', + ); + } + throw error; + } + const remainingAfterRefreshMs = localDeadlineAtMs - Date.now(); + if (remainingAfterRefreshMs <= 0) { + await this.rejectUnexecutedAssignment( + assignment, + 'Bridge assignment expired during credential refresh', + ); + throw new BridgeProtocolError( + 'Bridge assignment expired during credential refresh', + ); + } + if (signal != null && Boolean(signal.aborted)) { + throw signal.reason instanceof Error + ? signal.reason + : new DOMException('aborted', 'AbortError'); + } + const executionController = new AbortController(); + const credentialController = new AbortController(); + const abortExecution = (): void => { + executionController.abort(); + credentialController.abort(); + }; + signal?.addEventListener('abort', abortExecution, { once: true }); + const deadlineDelay = remainingAfterRefreshMs; + const deadlineTimer = setTimeout( + () => executionController.abort(), + deadlineDelay, + ); + if (this.lastRegisteredAtMs === 0) { + this.lastRegisteredAtMs = Date.now(); + } + const heartbeatController = new AbortController(); + let heartbeatError: unknown; + const heartbeat = this.maintainRegistration( + heartbeatController.signal, + ).catch((error) => { + heartbeatError = error; + executionController.abort(); + }); + const cancellationController = new AbortController(); + const cancellationWatcher = this.watchCancellation( + assignment, + executionController, + cancellationController.signal, + ); + let credentialMaintenanceError: unknown; + let credentialMaintenance: Promise | undefined; + let settlement: BridgeSettlement; + let ambiguousSandboxError: unknown; + let sandboxRejectedExecution = false; + let sandboxStarted = false; + try { + credentialMaintenance = this.maintainCredential( + assignment, + credentialController.signal, + serverClockOffsetMs, + signal, + ).catch((error) => { + credentialMaintenanceError = error; + executionController.abort(); + }); + const sandboxExecuteUrl = + `${this.sandboxEndpointFor(assignment)}/execute`; + const sandboxSessionId = this.sandboxSessionIdFor(assignment); + const headers = { + ...assignment.request.headers, + ...(sandboxSessionId + ? { 'X-Runtime-Session-Id': sandboxSessionId } + : {}), + }; + const sandboxRequestBody = JSON.stringify(assignment.request.body); + if (Date.now() >= localDeadlineAtMs) { + throw new BridgeProtocolError( + 'Bridge assignment expired before sandbox execution', + ); + } + sandboxStarted = true; + const response = await this.fetchImpl( + sandboxExecuteUrl, + { + method: 'POST', + headers: { + ...headers, + 'Content-Type': 'application/json', + }, + body: sandboxRequestBody, + signal: executionController.signal, + }, + ); + let payload: object = {}; + try { + payload = (await response.json()) as object; + } catch (error) { + if (response.ok) throw error; + } + if (credentialMaintenanceError != null) { + throw credentialMaintenanceError; + } + 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) ?? + `Sandbox rejected execution with HTTP ${response.status}`, + response.status, + ); + } + if (heartbeatError != null) throw heartbeatError; + settlement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, + status: 'fulfilled', + result: payload, + }; + } catch (error) { + if ( + assignment.runtimeSessionId != null && + sandboxStarted && + !sandboxRejectedExecution + ) { + ambiguousSandboxError = error; + } + settlement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, + status: 'rejected', + error: + (error instanceof Error + ? error.message + : 'Sandbox execution failed' + ).slice(0, MAX_SETTLEMENT_ERROR_LENGTH), + }; + } + + clearTimeout(deadlineTimer); + cancellationController.abort(); + await cancellationWatcher; + credentialController.abort(); + await credentialMaintenance; + try { + if (ambiguousSandboxError != null) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined after an ambiguous sandbox execution`, + ambiguousSandboxError, + ); + } + const knownCleanStatefulRejection = + assignment.runtimeSessionId != null && + settlement.status === 'rejected' && + (!sandboxStarted || sandboxRejectedExecution); + 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, + ), + ); + } finally { + recoveryHeartbeatController.abort(); + await recoveryHeartbeat; + } + } else { + await this.settleWithRetry( + assignment, + settlement, + localDeadlineAtMs, + 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 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)) { + return this.sandboxEndpoint; + } + return this.sandboxEndpoint.replace( + RUNTIME_SESSION_PLACEHOLDER, + encodeURIComponent(`assignment-${assignment.assignmentId}`), + ); + } + if ( + this.options.capabilities.statefulWorkspace !== true || + !this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER) + ) { + throw new BridgeProtocolError( + 'Stateful assignments require a sandbox endpoint template containing {runtimeSessionId}', + ); + } + return this.sandboxEndpoint.replace( + RUNTIME_SESSION_PLACEHOLDER, + encodeURIComponent(assignment.runtimeSessionId), + ); + } + + private async delay(ms: number, signal: AbortSignal): Promise { + await abortableDelay(ms, signal); + } + + private async maintainRegistration( + signal: AbortSignal, + retryTransient = false, + ): Promise { + while (!signal.aborted) { + const heartbeatIntervalMs = Math.max( + MIN_REGISTRATION_HEARTBEAT_MS, + Math.floor(this.registrationTtlMs / 2), + ); + await this.delay( + Math.max( + 0, + this.lastRegisteredAtMs + heartbeatIntervalMs - Date.now(), + ), + signal, + ); + if (signal.aborted) return; + 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; + } + } + + private assignmentUrl(assignment: BridgeAssignment, action: string): string { + return ( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + + `/assignments/${encodeURIComponent(assignment.assignmentId)}/${action}` + ); + } + + private async settleWithRetry( + assignment: BridgeAssignment, + settlement: BridgeSettlement, + 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 }); + const deadlineTimer = setTimeout( + () => settlementController.abort(), + Math.max(0, deadlineAtMs - Date.now()), + ); + let lastError: unknown; + try { + while (!settlementController.signal.aborted) { + try { + await this.request( + this.assignmentUrl(assignment, 'settle'), + settlement, + settlementController.signal, + ); + return; + } catch (error) { + lastError = error; + if (signal?.aborted) break; + if ( + error instanceof BridgeProtocolError && + error.status != null && + error.status < 500 && + error.status !== 408 && + error.status !== 429 + ) { + if ( + assignment.runtimeSessionId != null && + settlement.status === 'fulfilled' + ) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined after Code API rejected its fulfilled settlement`, + error, + ); + } + throw error; + } + const remainingMs = deadlineAtMs - Date.now(); + if (remainingMs <= 0) break; + await this.delay( + Math.min(SETTLEMENT_RETRY_DELAY_MS, remainingMs), + settlementController.signal, + ); + } + } + } finally { + clearTimeout(deadlineTimer); + signal?.removeEventListener('abort', abortSettlement); + } + if ( + assignment.runtimeSessionId != null && + settlement.status === 'fulfilled' + ) { + throw new BridgeWorkspaceQuarantinedError( + `Stateful workspace ${assignment.runtimeSessionId} was quarantined after ambiguous settlement delivery`, + lastError, + ); + } + if (lastError instanceof Error) throw lastError; + throw new BridgeProtocolError('Bridge settlement deadline expired'); + } + + private async watchCancellation( + assignment: BridgeAssignment, + executionController: AbortController, + signal: AbortSignal, + ): Promise { + while (!signal.aborted && !executionController.signal.aborted) { + await 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'), + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + }, + pollController.signal, + ); + if (response.cancelled) { + executionController.abort(); + return; + } + } catch (error) { + if (signal.aborted) return; + if (error instanceof BridgeProtocolError && error.status === 404) { + executionController.abort(); + return; + } + } finally { + clearTimeout(timeout); + signal.removeEventListener('abort', abortPoll); + } + } + } + + private async request( + url: string, + body: object, + signal?: AbortSignal, + ): Promise { + const requestBody = JSON.stringify(body); + const response = await this.fetchImpl(url, { + method: 'POST', + headers: { + ...this.authorizationHeaders(url, requestBody), + 'Content-Type': 'application/json', + }, + body: requestBody, + signal, + }); + 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(errorPayload) ?? + `Bridge request failed with HTTP ${response.status}`, + response.status, + errorCode(errorPayload), + ); + } + return payload as T; + } + + private authorizationHeaders( + url: string, + body: string, + ): Record { + const identity = this.options.identity; + if (identity == null) { + return { Authorization: `Bearer ${this.options.token}` }; + } + const timestamp = new Date().toISOString(); + const nonce = randomBytes(18).toString('base64url'); + const proof = { + credential: identity.credential, + method: 'POST', + path: new URL(url).pathname, + timestamp, + nonce, + body, + }; + return { + Authorization: `Bridge ${identity.credential}`, + 'X-LibreChat-Code-Timestamp': timestamp, + 'X-LibreChat-Code-Nonce': nonce, + 'X-LibreChat-Code-Signature': signBridgeRequest( + identity.privateKey, + proof, + ), + }; + } + + 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/packages/code/tsconfig.json b/packages/code/tsconfig.json new file mode 100644 index 00000000..6c7ea362 --- /dev/null +++ b/packages/code/tsconfig.json @@ -0,0 +1,15 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "NodeNext", + "moduleResolution": "NodeNext", + "strict": true, + "declaration": true, + "sourceMap": true, + "outDir": "dist", + "rootDir": "src", + "skipLibCheck": true, + "types": ["node"] + }, + "include": ["src/**/*.ts"] +} diff --git a/service/Dockerfile b/service/Dockerfile index 762790dc..00680111 100644 --- a/service/Dockerfile +++ b/service/Dockerfile @@ -18,6 +18,7 @@ COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY service/scripts ./scripts COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ RUN bun build ./src/file-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' RUN bun build ./src/api-server.ts --minify --outdir .build-api --target bun --external '@opentelemetry/*' @@ -66,6 +67,7 @@ ENV NODE_ENV=development COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src COPY shared /shared +COPY packages/code/src /packages/code/src COPY service/tsconfig.json ./ EXPOSE 3000 9230 CMD ["bun", "run", "--watch", "src/file-server.ts"] diff --git a/service/Dockerfile.api b/service/Dockerfile.api index 419bdc95..2921401e 100644 --- a/service/Dockerfile.api +++ b/service/Dockerfile.api @@ -18,8 +18,10 @@ 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 packages/code /packages/code COPY service/tsconfig.json ./ RUN bun build ./src/api-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' RUN bun build ./scripts/rehydrate-session-cache.ts --minify --outdir .build-migrations --target bun --external '@opentelemetry/*' @@ -46,7 +48,9 @@ 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 packages/code /packages/code COPY service/tsconfig.json ./ EXPOSE 3112 9230 CMD ["bun", "run", "--watch", "src/api-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/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/Dockerfile.worker b/service/Dockerfile.worker index 9d1a4322..e99c16c4 100644 --- a/service/Dockerfile.worker +++ b/service/Dockerfile.worker @@ -19,6 +19,7 @@ RUN cd /temp/prod && bun install --frozen-lockfile --production FROM base AS builder COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ RUN bun build ./src/worker-server.ts --minify --outdir .build --target bun --external '@opentelemetry/*' @@ -43,6 +44,7 @@ FROM base AS development ENV NODE_ENV=development COPY --from=install /temp/dev/node_modules node_modules COPY service/src ./src +COPY packages/code/src /packages/code/src COPY shared /shared COPY service/tsconfig.json ./ EXPOSE 3113 9230 diff --git a/service/rollup.config.js b/service/rollup.config.js index 2400f72d..059c71d1 100644 --- a/service/rollup.config.js +++ b/service/rollup.config.js @@ -38,7 +38,12 @@ 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', + '../packages/code/src/identity.ts', + ], sourceMap: true, declaration: false, declarationMap: false, diff --git a/service/src/api-server.ts b/service/src/api-server.ts index 78689826..1e4634a4 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'; 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/auth/librechat-jwt.test.ts b/service/src/auth/librechat-jwt.test.ts index 2030b2e7..d20124f5 100644 --- a/service/src/auth/librechat-jwt.test.ts +++ b/service/src/auth/librechat-jwt.test.ts @@ -48,6 +48,7 @@ type JwtClaims = { chc_user_id?: string; auth_context_hash?: string; plan_id?: string; + code_worker_id?: string; }; const originalEnv = new Map(); @@ -75,6 +76,7 @@ function baseClaims(overrides: Partial = {}): JwtClaims { external_user_id: 'chc_123', auth_context_hash: 'hash_123', plan_id: 'prod_plan_123', + code_worker_id: 'code-user_123', ...overrides, }; } @@ -166,6 +168,7 @@ describe('LibreChat JWT auth provider', () => { principalSource: 'openid_reuse', authContextHash: 'hash_123', planId: 'prod_plan_123', + codeWorkerId: 'code-user_123', }); }); diff --git a/service/src/auth/librechat-jwt.ts b/service/src/auth/librechat-jwt.ts index 1e7e8079..e249ce0c 100644 --- a/service/src/auth/librechat-jwt.ts +++ b/service/src/auth/librechat-jwt.ts @@ -39,6 +39,7 @@ interface LibreChatJwtClaims { chc_user_id?: string; // leak-check:allow auth_context_hash?: string; plan_id?: string; + code_worker_id?: string; } interface PublicKeyEntry { @@ -394,6 +395,7 @@ function validateClaims(claims: LibreChatJwtClaims, config: VerificationConfig): const nbf = assertNumericDate(claims.nbf, 'nbf'); const exp = assertNumericDate(claims.exp, 'exp'); const planId = optionalString(claims.plan_id, 'plan_id'); + const codeWorkerId = optionalString(claims.code_worker_id, 'code_worker_id'); const principalSource = assertPrincipalSource(claims.principal_source); const authContextHash = assertString(claims.auth_context_hash, 'auth_context_hash'); @@ -433,6 +435,7 @@ function validateClaims(claims: LibreChatJwtClaims, config: VerificationConfig): principalSource, authContextHash, planId, + codeWorkerId, }; } diff --git a/service/src/auth/principal.ts b/service/src/auth/principal.ts index 94b785f3..615a0ff2 100644 --- a/service/src/auth/principal.ts +++ b/service/src/auth/principal.ts @@ -13,6 +13,7 @@ export type CodeApiPrincipal = { authContextHash?: string; credentialId?: string; planId?: string; + codeWorkerId?: string; }; export function applyPrincipal(req: t.AuthenticatedRequest, principal: CodeApiPrincipal): void { diff --git a/service/src/bridge/index.ts b/service/src/bridge/index.ts new file mode 100644 index 00000000..08ef1b07 --- /dev/null +++ b/service/src/bridge/index.ts @@ -0,0 +1,17 @@ +import { connection } from '../queue'; +import { env } from '../config'; +import { RedisBridgePairingStore } from './pairing'; +import { createBridgeRouter } from './router'; +import { RedisBridgeStore } from './store'; + +export const bridgeStore = new RedisBridgeStore(connection); +export const bridgePairings = new RedisBridgePairingStore(connection); + +export default createBridgeRouter({ + store: bridgeStore, + pairings: bridgePairings, + authMode: env.BRIDGE_AUTH_MODE, + adminToken: env.BRIDGE_TOKEN, + configuredWorkerId: env.BRIDGE_WORKER_ID, + allowDynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, +}); diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts new file mode 100644 index 00000000..36b2ea6d --- /dev/null +++ b/service/src/bridge/pairing.test.ts @@ -0,0 +1,1144 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { createHash } from 'crypto'; +import RedisMock from 'ioredis-mock'; + +import type Redis from 'ioredis'; + +import { + createBridgeIdentity, + signBridgeRequest, +} from '../../../packages/code/src/identity'; +import { RedisBridgePairingStore } from './pairing'; +import { RedisBridgeStore } from './store'; + +const redis = new RedisMock() as unknown as Redis; +const pairings = new RedisBridgePairingStore(redis); +const store = new RedisBridgeStore(redis); + +afterEach(async () => { + await redis.flushall(); +}); + +describe('RedisBridgePairingStore', () => { + test('preserves a tenant and generic principal binding across credential rotation', async () => { + const identity = createBridgeIdentity(); + const binding = { + tenantId: 'tenant-1', + principal: { type: 'group' as const, id: 'engineering' }, + }; + const pairing = await pairings.issue('vm-bound', binding); + const issued = await pairings.redeem({ + workerId: 'vm-bound', + code: pairing.code, + publicKey: identity.publicKey, + }); + const requestFor = ( + credential: string, + nonce: string, + ): Parameters[0] => { + const proof = { + credential, + method: 'POST', + path: '/v1/bridge/workers/vm-bound/lease', + timestamp: new Date().toISOString(), + nonce, + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + return { + ...proof, + workerId: 'vm-bound', + signature: signBridgeRequest(identity.privateKey, proof), + }; + }; + + const originalAuthorization = await pairings.authorize( + requestFor(issued.credential, 'original-bound-worker-proof'), + ); + const rotated = await pairings.rotate('vm-bound'); + + const rotatedAuthorization = await pairings.authorize( + requestFor(rotated.credential, 'bound-worker-proof'), + ); + expect(rotatedAuthorization).toMatchObject({ workerId: 'vm-bound', binding }); + expect(typeof originalAuthorization.identityId).toBe('string'); + expect(rotatedAuthorization.identityId).toBe( + originalAuthorization.identityId, + ); + await expect( + pairings.authorize(requestFor(issued.credential, 'superseded-bound-proof')), + ).resolves.toMatchObject({ + workerId: 'vm-bound', + identityId: originalAuthorization.identityId, + }); + }); + + test('preserves a legacy unmarked identity across its first rotation', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('legacy-vm'); + const issued = await pairings.redeem({ + workerId: 'legacy-vm', + code: pairing.code, + publicKey: identity.publicKey, + }); + const issuedDigest = createHash('sha256') + .update(issued.credential) + .digest('hex'); + const credentialKey = `codeapi:bridge:v1:credential:${issuedDigest}`; + const stored = JSON.parse((await redis.get(credentialKey)) ?? '{}') as { + identityId?: string; + }; + delete stored.identityId; + await redis.set(credentialKey, JSON.stringify(stored), 'EX', 300); + + const rotated = await pairings.rotate('legacy-vm'); + const proof = { + credential: rotated.credential, + method: 'POST', + path: '/v1/bridge/workers/legacy-vm/lease', + timestamp: new Date().toISOString(), + nonce: 'legacy-rotation-proof', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + const authorization = await pairings.authorize({ + ...proof, + workerId: 'legacy-vm', + signature: signBridgeRequest(identity.privateKey, proof), + }); + + expect(authorization.identityId).toBeUndefined(); + }); + + test('redeems a pairing code exactly once for the intended worker identity', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + + const credential = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + + expect(credential.workerId).toBe('vm-1'); + expect(credential.credential.length).toBeGreaterThanOrEqual(32); + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + }); + + test('preserves a pairing code after public-key validation fails', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-public-key-retry'); + + await expect( + pairings.redeem({ + workerId: 'vm-public-key-retry', + code: pairing.code, + publicKey: 'not-a-public-key', + }), + ).rejects.toMatchObject({ code: 'PUBLIC_KEY_INVALID' }); + + await expect( + pairings.redeem({ + workerId: 'vm-public-key-retry', + code: pairing.code, + publicKey: identity.publicKey, + }), + ).resolves.toMatchObject({ workerId: 'vm-public-key-retry' }); + }); + + test('only the newest pairing code can rebind a worker identity', async () => { + const identity = createBridgeIdentity(); + const older = await pairings.issue('vm-1', { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }); + const newer = await pairings.issue('vm-1', { + tenantId: 'tenant-b', + principal: { type: 'user', id: 'user-b' }, + }); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: older.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: newer.code, + publicKey: identity.publicKey, + }), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('does not let a paused redemption overwrite a newer pairing identity', async () => { + const firstIdentity = createBridgeIdentity(); + const secondIdentity = createBridgeIdentity(); + const firstPairing = await pairings.issue('vm-race', { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }); + const originalEval = redis.eval.bind(redis); + let releaseFirst!: () => void; + let firstRedeemed!: () => void; + const firstRedeemedPromise = new Promise((resolve) => { + firstRedeemed = resolve; + }); + const releaseFirstPromise = new Promise((resolve) => { + releaseFirst = resolve; + }); + let paused = false; + redis.eval = (async (script: string, ...args: unknown[]) => { + if (!paused && script.includes('if pairing ~= ARGV[1]')) { + paused = true; + firstRedeemed(); + await releaseFirstPromise; + } + return await (originalEval as (...evalArgs: unknown[]) => Promise)( + script, + ...args, + ); + }) as typeof redis.eval; + + try { + const staleRedemption = pairings.redeem({ + workerId: 'vm-race', + code: firstPairing.code, + publicKey: firstIdentity.publicKey, + }); + await firstRedeemedPromise; + const secondPairing = await pairings.issue('vm-race', { + tenantId: 'tenant-b', + principal: { type: 'user', id: 'user-b' }, + }); + const current = await pairings.redeem({ + workerId: 'vm-race', + code: secondPairing.code, + publicKey: secondIdentity.publicKey, + }); + releaseFirst(); + + await expect(staleRedemption).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + const proof = { + credential: current.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-race/lease', + timestamp: new Date().toISOString(), + nonce: 'current-race-proof', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-race', + signature: signBridgeRequest(secondIdentity.privateKey, proof), + }), + ).resolves.toMatchObject({ + binding: { + tenantId: 'tenant-b', + principal: { type: 'user', id: 'user-b' }, + }, + }); + } finally { + redis.eval = originalEval as typeof redis.eval; + releaseFirst(); + } + }); + + test('authorizes a credential only with proof from its worker key', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce: 'request-nonce-1', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('rejects replay of an already accepted worker proof', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce: 'single-use-nonce', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + const request = { + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }; + + await pairings.authorize(request); + + await expect(pairings.authorize(request)).rejects.toMatchObject({ + code: 'PROOF_REPLAYED', + }); + }); + + test('rejects a correctly signed proof outside the clock window', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date(Date.now() - 5 * 60_000).toISOString(), + nonce: 'stale-request-nonce', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }), + ).rejects.toMatchObject({ code: 'PROOF_INVALID' }); + }); + + test('revocation immediately invalidates the active worker credential', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: new Date().toISOString(), + nonce: 'post-revocation-request', + body: JSON.stringify({ protocolVersion: 1, workerId: 'vm-1' }), + }; + await store.register({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + + await pairings.revoke('vm-1'); + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + expect(await redis.get('codeapi:bridge:v1:worker:vm-1')).toBeNull(); + expect( + await redis.get('codeapi:bridge:v1:worker:vm-1:incarnation'), + ).toBeNull(); + await expect( + store.register({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_FENCED' }); + }); + + test('revocation fences a registration authorized before the revoke', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: new Date().toISOString(), + nonce: 'registration-revoke-race', + body: JSON.stringify({ protocolVersion: 1, workerId: 'vm-1' }), + }; + const authorization = await pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }); + + await pairings.revoke('vm-1'); + + await expect( + store.register( + { + protocolVersion: 1, + workerId: 'vm-1', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }, + authorization, + ), + ).rejects.toMatchObject({ code: 'WORKER_FENCED' }); + expect(await redis.get('codeapi:bridge:v1:worker:vm-1')).toBeNull(); + expect( + await redis.get('codeapi:bridge:v1:worker:vm-1:incarnation'), + ).toBeNull(); + }); + + test('revocation invalidates pairing codes issued before the revoke', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + + await pairings.revoke('vm-1'); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + }); + + test('revocation invalidates an unredeemed pairing code', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + + await pairings.revoke('vm-1'); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + }); + + test('revocation removes pairing codes issued by a pre-fence replica', async () => { + const legacyCode = 'legacy-pairing-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ workerId: 'vm-1', expiresAt: new Date(Date.now() + 60_000).toISOString() }), + 'EX', + 60, + ); + + await pairings.revoke('vm-1'); + + await expect(redis.get(legacyKey)).resolves.toBeNull(); + }); + + test('retries legacy cleanup after a transient scan failure', async () => { + const legacyCode = 'retryable-legacy-pairing-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-scan-retry', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + const scan = redis.scan.bind(redis); + let failScan = true; + redis.scan = (async (...args: Parameters) => { + if (failScan) { + failScan = false; + throw new Error('transient scan failure'); + } + return scan(...args); + }) as Redis['scan']; + + await expect(pairings.revoke('vm-scan-retry')).rejects.toThrow( + 'transient scan failure', + ); + redis.scan = scan; + await pairings.revoke('vm-scan-retry'); + + await expect(redis.get(legacyKey)).resolves.toBeNull(); + }); + + test('retries a claimed cleanup after the migration deadline', async () => { + const legacyCode = 'post-deadline-retry-pairing-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-post-deadline-retry', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + await redis.set( + 'codeapi:bridge:v1:migration:legacy-pairing-scan-until', + String(Date.now() + 15), + ); + const scan = redis.scan.bind(redis); + let failScan = true; + redis.scan = (async (...args: Parameters) => { + if (failScan) { + failScan = false; + await new Promise((resolve) => setTimeout(resolve, 30)); + throw new Error('scan failed after deadline'); + } + return scan(...args); + }) as Redis['scan']; + + await expect(pairings.revoke('vm-post-deadline-retry')).rejects.toThrow( + 'scan failed after deadline', + ); + await pairings.revoke('vm-post-deadline-retry'); + + redis.scan = scan; + await expect(redis.get(legacyKey)).resolves.toBeNull(); + }); + + test('renews the cleanup claim while a shared-keyspace scan is in flight', async () => { + const store = new RedisBridgePairingStore(redis, 600, 300, 30); + const legacyCode = 'renewed-claim-pairing-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-renewed-claim', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + const scan = redis.scan.bind(redis); + let scanCalls = 0; + let renewCalls = 0; + let markScanStarted = () => {}; + const scanStarted = new Promise((resolve) => { + markScanStarted = resolve; + }); + redis.scan = (async (...args: Parameters) => { + scanCalls += 1; + if (scanCalls === 1) { + markScanStarted(); + await new Promise((resolve) => setTimeout(resolve, 80)); + } + return scan(...args); + }) as Redis['scan']; + const originalEval = redis.eval.bind(redis); + redis.eval = (async (script: string, ...args: unknown[]) => { + if (script.includes("redis.call('PEXPIRE'")) renewCalls += 1; + return await (originalEval as (...evalArgs: unknown[]) => Promise)( + script, + ...args, + ); + }) as typeof redis.eval; + + try { + const first = store.revoke('vm-renewed-claim'); + await scanStarted; + await new Promise((resolve) => setTimeout(resolve, 45)); + expect(renewCalls).toBeGreaterThan(0); + await first; + + expect(scanCalls).toBe(1); + await expect(redis.get(legacyKey)).resolves.toBeNull(); + } finally { + redis.scan = scan; + redis.eval = originalEval as typeof redis.eval; + } + }); + + test('rescans an ambiguous marker written by the preceding build', async () => { + const legacyCode = 'ambiguous-predecessor-marker-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + const workerId = 'vm-ambiguous-predecessor-marker'; + await redis.set( + legacyKey, + JSON.stringify({ + workerId, + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + await redis.set( + `codeapi:bridge:v1:migration:legacy-pairing-scanned:${workerId}`, + 'predecessor-random-token', + 'PX', + 60_000, + ); + + await pairings.revoke(workerId); + + await expect(redis.get(legacyKey)).resolves.toBeNull(); + }); + + test('starts the advertised pairing lifetime after legacy cleanup', async () => { + const store = new RedisBridgePairingStore(redis, 60); + const scan = redis.scan.bind(redis); + const originalNow = Date.now; + let now = originalNow(); + Date.now = () => now; + redis.scan = (async (...args: Parameters) => { + const result = await scan(...args); + now += 10; + return result; + }) as Redis['scan']; + + try { + const pairing = await store.issue('vm-post-cleanup-expiry'); + expect(Date.parse(pairing.expiresAt) - Date.now()).toBe(60_000); + } finally { + Date.now = originalNow; + redis.scan = scan; + } + }); + + test('waits for a failed in-progress cleanup and confirms legacy removal itself', async () => { + const legacyCode = 'overlapping-legacy-pairing-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-overlapping-cleanup', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + const scan = redis.scan.bind(redis); + let releaseFirstScan = () => {}; + const firstScanGate = new Promise((resolve) => { + releaseFirstScan = resolve; + }); + let markFirstScanStarted = () => {}; + const firstScanStarted = new Promise((resolve) => { + markFirstScanStarted = resolve; + }); + let scanCalls = 0; + redis.scan = (async (...args: Parameters) => { + scanCalls += 1; + if (scanCalls === 1) { + markFirstScanStarted(); + await firstScanGate; + throw new Error('interrupted claimed scan'); + } + return scan(...args); + }) as Redis['scan']; + + const interrupted = pairings.revoke('vm-overlapping-cleanup'); + await firstScanStarted; + const overlapping = pairings.revoke('vm-overlapping-cleanup'); + let overlappingSettled = false; + void overlapping.then( + () => { + overlappingSettled = true; + }, + () => { + overlappingSettled = true; + }, + ); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(overlappingSettled).toBe(false); + releaseFirstScan(); + + await expect(interrupted).rejects.toThrow('interrupted claimed scan'); + await expect(overlapping).resolves.toBeUndefined(); + redis.scan = scan; + await expect(redis.get(legacyKey)).resolves.toBeNull(); + }); + + test('legacy cleanup does not delete generation-fenced pairings', async () => { + const fencedCode = 'concurrent-generation-pairing'; + const fencedKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(fencedCode) + .digest('hex')}`; + await redis.set( + fencedKey, + JSON.stringify({ + workerId: 'vm-generation-fenced', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + generation: 'new-generation', + }), + 'EX', + 60, + ); + + await pairings.revoke('vm-generation-fenced'); + + await expect(redis.get(fencedKey)).resolves.not.toBeNull(); + }); + + test('reopens legacy cleanup after a rollback outlives the prior scan window', async () => { + const legacyCode = 'later-rollback-pairing-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + const deadlineKey = 'codeapi:bridge:v1:migration:legacy-pairing-scan-until'; + await redis.set(deadlineKey, '0'); + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-later-rollback', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + await redis.set( + 'codeapi:bridge:v1:pairing-index:vm-later-rollback', + legacyKey, + 'EX', + 60, + ); + await redis.set( + 'codeapi:bridge:v1:migration:legacy-pairing-scanned:vm-later-rollback', + 'done', + ); + + await pairings.revoke('vm-later-rollback'); + + await expect(redis.get(legacyKey)).resolves.toBeNull(); + await expect(redis.get(deadlineKey)).resolves.toBe('0'); + }); + + test('does not restart an expired migration window without rollback evidence', async () => { + const legacyCode = 'unindexed-post-migration-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + const deadlineKey = 'codeapi:bridge:v1:migration:legacy-pairing-scan-until'; + await redis.set(deadlineKey, '0'); + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-no-rollback-signal', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + + await pairings.revoke('vm-no-rollback-signal'); + + await expect(redis.get(legacyKey)).resolves.not.toBeNull(); + await expect(redis.get(deadlineKey)).resolves.toBe('0'); + }); + + test('restarts legacy cleanup once for an explicit rollback epoch', async () => { + const legacyCode = 'unindexed-rollback-epoch-code'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + const deadlineKey = 'codeapi:bridge:v1:migration:legacy-pairing-scan-until'; + await redis.set(deadlineKey, '0'); + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-rollback-epoch', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + await redis.set( + 'codeapi:bridge:v1:migration:legacy-pairing-scanned:vm-rollback-epoch', + 'done', + ); + const rollbackAwarePairings = new RedisBridgePairingStore( + redis, + 600, + 300, + 5_000, + 'rollback-epoch-1', + ); + + await rollbackAwarePairings.revoke('vm-rollback-epoch'); + + await expect(redis.get(legacyKey)).resolves.toBeNull(); + await expect(redis.get(deadlineKey)).resolves.toBe('0'); + const epochHash = createHash('sha256') + .update('rollback-epoch-1') + .digest('hex'); + const epochStateKey = + `codeapi:bridge:v1:migration:legacy-pairing-scanned:` + + `vm-rollback-epoch:${epochHash}`; + await expect(redis.get(epochStateKey)).resolves.toBe('done'); + await expect(redis.pttl(epochStateKey)).resolves.toBe(-1); + }); + + test('cleans a revoked legacy code before rollback-epoch redemption', async () => { + const identity = createBridgeIdentity(); + const workerId = 'vm-rollback-redeem'; + const legacyCode = 'revoked-rollback-pairing'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set('codeapi:bridge:v1:migration:legacy-pairing-scan-until', '0'); + await redis.set( + legacyKey, + JSON.stringify({ + workerId, + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + const rollbackAwarePairings = new RedisBridgePairingStore( + redis, + 600, + 300, + 5_000, + 'rollback-epoch-redeem', + ); + + await expect( + rollbackAwarePairings.redeem({ + workerId, + code: legacyCode, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + await expect(redis.get(legacyKey)).resolves.toBeNull(); + }); + + test('does not scan for a nonexistent rollback-epoch pairing code', async () => { + const identity = createBridgeIdentity(); + const scan = redis.scan.bind(redis); + let scanCalls = 0; + redis.scan = (async (...args: Parameters) => { + scanCalls += 1; + return scan(...args); + }) as Redis['scan']; + const rollbackAwarePairings = new RedisBridgePairingStore( + redis, + 600, + 300, + 5_000, + 'rollback-epoch-missing-code', + ); + + try { + await expect( + rollbackAwarePairings.redeem({ + workerId: 'attacker-chosen-worker', + code: 'nonexistent-code', + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + expect(scanCalls).toBe(0); + await expect( + redis.keys( + 'codeapi:bridge:v1:migration:legacy-pairing-scanned:attacker-chosen-worker*', + ), + ).resolves.toEqual([]); + } finally { + redis.scan = scan; + } + }); + + test('redeems an unrevoked pairing code issued by a pre-fence replica', async () => { + const identity = createBridgeIdentity(); + const legacyCode = 'unrevoked-legacy-pairing'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ workerId: 'vm-1', expiresAt: new Date(Date.now() + 60_000).toISOString() }), + 'EX', + 60, + ); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: legacyCode, + publicKey: identity.publicKey, + }), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('redeems an indexed legacy code issued after rollback', async () => { + const identity = createBridgeIdentity(); + await pairings.issue('vm-rollback'); + const legacyCode = 'rollback-issued-legacy-pairing'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ + workerId: 'vm-rollback', + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + 'EX', + 60, + ); + await redis.set( + 'codeapi:bridge:v1:pairing-index:vm-rollback', + legacyKey, + 'EX', + 60, + ); + + await expect( + pairings.redeem({ + workerId: 'vm-rollback', + code: legacyCode, + publicKey: identity.publicKey, + }), + ).resolves.toMatchObject({ workerId: 'vm-rollback' }); + }); + + test('replacement invalidates a pairing code issued by a pre-fence replica', async () => { + const identity = createBridgeIdentity(); + const legacyCode = 'replaced-legacy-pairing'; + const legacyKey = `codeapi:bridge:v1:pairing:${createHash('sha256') + .update(legacyCode) + .digest('hex')}`; + await redis.set( + legacyKey, + JSON.stringify({ workerId: 'vm-1', expiresAt: new Date(Date.now() + 60_000).toISOString() }), + 'EX', + 60, + ); + + await pairings.issue('vm-1'); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: legacyCode, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + }); + + test('issuing a replacement invalidates the prior unredeemed pairing code', async () => { + const identity = createBridgeIdentity(); + const first = await pairings.issue('vm-1'); + const replacement = await pairings.issue('vm-1'); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: first.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: replacement.code, + publicKey: identity.publicKey, + }), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('rotation retains the prior same-identity credential for recovery', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const original = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + + const rotated = await pairings.rotate('vm-1'); + + const proofFor = ( + credential: string, + nonce: string, + ): Parameters[0] => { + const proof = { + credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce, + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + return { + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }; + }; + + await expect( + pairings.authorize(proofFor(original.credential, 'old-credential')), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + await expect( + pairings.authorize(proofFor(rotated.credential, 'new-credential')), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('recovers when a refresh response is lost after the server commits it', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const original = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proofFor = ( + credential: string, + nonce: string, + ): Parameters[0] => { + const proof = { + credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/credentials/refresh', + timestamp: new Date().toISOString(), + nonce, + body: JSON.stringify({ protocolVersion: 1 }), + }; + return { + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }; + }; + + await pairings.rotate('vm-1'); + const retryAuthorization = await pairings.authorize( + proofFor(original.credential, 'refresh-response-lost'), + ); + const recovered = await pairings.rotate( + 'vm-1', + retryAuthorization.credentialId, + ); + + await expect( + pairings.authorize(proofFor(recovered.credential, 'refresh-recovered')), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('rejects a stale credential refresh after the worker is paired again', async () => { + const originalIdentity = createBridgeIdentity(); + const originalPairing = await pairings.issue('vm-1'); + const original = await pairings.redeem({ + workerId: 'vm-1', + code: originalPairing.code, + publicKey: originalIdentity.publicKey, + }); + const proof = { + credential: original.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/credentials/refresh', + timestamp: new Date().toISOString(), + nonce: 'authorized-before-repairing', + body: JSON.stringify({ protocolVersion: 1 }), + }; + const staleAuthorization = await pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(originalIdentity.privateKey, proof), + }); + + const replacementIdentity = createBridgeIdentity(); + const replacementPairing = await pairings.issue('vm-1'); + await pairings.redeem({ + workerId: 'vm-1', + code: replacementPairing.code, + publicKey: replacementIdentity.publicKey, + }); + + await expect( + pairings.rotate('vm-1', staleAuthorization.credentialId), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + }); + + test('repairing a worker invalidates its previously paired credential', async () => { + const firstIdentity = createBridgeIdentity(); + const firstPairing = await pairings.issue('vm-1'); + const first = await pairings.redeem({ + workerId: 'vm-1', + code: firstPairing.code, + publicKey: firstIdentity.publicKey, + }); + const nextIdentity = createBridgeIdentity(); + const nextPairing = await pairings.issue('vm-1'); + await pairings.redeem({ + workerId: 'vm-1', + code: nextPairing.code, + publicKey: nextIdentity.publicKey, + }); + const proof = { + credential: first.credential, + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: new Date().toISOString(), + nonce: 'superseded-pairing', + body: JSON.stringify({ protocolVersion: 1, workerId: 'vm-1' }), + }; + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(firstIdentity.privateKey, proof), + }), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + }); +}); diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts new file mode 100644 index 00000000..74c75b7c --- /dev/null +++ b/service/src/bridge/pairing.ts @@ -0,0 +1,734 @@ +import { + createHash, + createPublicKey, + randomBytes, +} from 'crypto'; + +import type Redis from 'ioredis'; + +import { verifyBridgeRequest } from '../../../packages/code/src/identity'; + +const PREFIX = 'codeapi:bridge:v1'; +const DEFAULT_PAIRING_TTL_SECONDS = 10 * 60; +const DEFAULT_CREDENTIAL_TTL_SECONDS = 15 * 60; +const PROOF_NONCE_TTL_SECONDS = 2 * 60; +const PROOF_CLOCK_SKEW_MS = 60_000; +const LEGACY_SCAN_CLAIM_TTL_MS = 5_000; +const LEGACY_SCAN_POLL_INTERVAL_MS = 25; +const LEGACY_SCAN_PENDING = 'pending'; +const LEGACY_SCAN_COMPLETE = 'done'; +const ISSUE_PAIRING_SCRIPT = ` +local previous = redis.call('GET', KEYS[1]) +if previous then + redis.call('DEL', previous) +end +redis.call('SET', KEYS[1], KEYS[2], 'EX', ARGV[2]) +redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2]) +return 1 +`; +const REDEEM_PAIRING_SCRIPT = ` +local pairing = redis.call('GET', KEYS[1]) +if pairing ~= ARGV[1] then + return 0 +end +local generation = redis.call('GET', KEYS[2]) +if ARGV[2] == '' then + if generation and redis.call('GET', KEYS[5]) ~= KEYS[1] then + redis.call('DEL', KEYS[1]) + return 0 + end +elseif (generation or '0') ~= ARGV[2] then + redis.call('DEL', KEYS[1]) + return 0 +end +redis.call('DEL', KEYS[1]) +if redis.call('GET', KEYS[5]) == KEYS[1] then + redis.call('DEL', KEYS[5]) +end +redis.call('SET', KEYS[3], ARGV[3], 'EX', ARGV[4]) +redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +redis.call('SET', KEYS[6], ARGV[6], 'EX', ARGV[4]) +return 1 +`; +const ROTATE_CREDENTIAL_SCRIPT = ` +local activeDigest = redis.call('GET', KEYS[1]) +local previous = redis.call('GET', KEYS[2]) +if not activeDigest or not previous then + return 0 +end +if activeDigest ~= ARGV[1] then + if ARGV[5] == '' or redis.call('GET', KEYS[4]) ~= ARGV[5] then + return 0 + end +end +redis.call('SET', KEYS[3], ARGV[3], 'EX', ARGV[4]) +redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[4]) +if ARGV[5] ~= '' then + redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +else + redis.call('DEL', KEYS[4]) +end +return 1 +`; +const REVOKE_PAIRING_SCRIPT = ` +local indexed = redis.call('GET', KEYS[1]) +local credential = redis.call('GET', KEYS[3]) +local activeIncarnation = redis.call('GET', KEYS[6]) +redis.call('INCR', KEYS[2]) +if indexed then + redis.call('DEL', indexed) +end +if credential then + redis.call('DEL', KEYS[3]) + redis.call('DEL', ARGV[1] .. credential) +end +redis.call('DEL', KEYS[1], KEYS[3], KEYS[4], KEYS[5], KEYS[6]) +if activeIncarnation then + redis.call('SET', ARGV[2] .. activeIncarnation .. ':fenced', '1') +end +return 1 +`; +const RELEASE_LEGACY_SCAN_CLAIM_SCRIPT = ` +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +`; +const NORMALIZE_LEGACY_SCAN_STATE_SCRIPT = ` +if redis.call('GET', KEYS[1]) == ARGV[1] then + redis.call('SET', KEYS[1], ARGV[2]) + return 1 +end +return 0 +`; +const RENEW_LEGACY_SCAN_CLAIM_SCRIPT = ` +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('PEXPIRE', KEYS[1], ARGV[2]) +end +return 0 +`; +const COMPLETE_LEGACY_SCAN_CLAIM_SCRIPT = ` +if redis.call('GET', KEYS[2]) == ARGV[1] then + local remaining = tonumber(ARGV[3]) + if ARGV[4] == '1' then + redis.call('SET', KEYS[1], ARGV[2]) + elseif remaining > 0 then + redis.call('SET', KEYS[1], ARGV[2], 'PX', remaining) + else + redis.call('DEL', KEYS[1]) + end + redis.call('DEL', KEYS[2]) + return 1 +end +return 0 +`; + +export type BridgePrincipalType = 'deployment' | 'tenant' | 'user' | 'role' | 'group'; + +export interface BridgeWorkerBinding { + tenantId: string; + principal: { + type: BridgePrincipalType; + id: string; + }; +} + +interface StoredPairing { + workerId: string; + expiresAt: string; + generation?: number; + binding?: BridgeWorkerBinding; +} + +interface StoredCredential { + workerId: string; + /** Stable across refreshes; replaced only when the worker is paired again. */ + identityId?: string; + publicKey: string; + expiresAt: string; + binding?: BridgeWorkerBinding; +} + +export interface BridgePairing { + workerId: string; + code: string; + expiresAt: string; +} + +export interface BridgeWorkerCredential { + workerId: string; + credential: string; + expiresAt: string; +} + +export class BridgePairingError extends Error { + constructor( + public readonly code: + | 'PAIRING_INVALID' + | 'PUBLIC_KEY_INVALID' + | 'CREDENTIAL_INVALID' + | 'PROOF_INVALID' + | 'PROOF_REPLAYED', + message: string, + ) { + super(message); + this.name = 'BridgePairingError'; + } +} + +function digest(value: string): string { + return createHash('sha256').update(value).digest('hex'); +} + +function pairingKey(code: string): string { + return `${PREFIX}:pairing:${digest(code)}`; +} + +function credentialDigestKey(credentialDigest: string): string { + return `${PREFIX}:credential:${credentialDigest}`; +} + +function workerIdentityKey(workerId: string): string { + return `${PREFIX}:identity:${workerId}`; +} + +function workerStableIdentityKey(workerId: string): string { + return `${PREFIX}:stable-identity:${workerId}`; +} + +function workerPairingGenerationKey(workerId: string): string { + return `${PREFIX}:pairing-generation:${workerId}`; +} + +function workerPairingIndexKey(workerId: string): string { + return `${PREFIX}:pairing-index:${workerId}`; +} + +function legacyPairingScanDeadlineKey(): string { + return `${PREFIX}:migration:legacy-pairing-scan-until`; +} + +function legacyPairingWorkerScanKey( + workerId: string, + rollbackEpoch?: string, +): string { + const epoch = rollbackEpoch?.trim(); + const epochSuffix = epoch ? `:${digest(epoch)}` : ''; + return `${PREFIX}:migration:legacy-pairing-scanned:${workerId}${epochSuffix}`; +} + +function proofNonceKey(credential: string, nonce: string): string { + return `${PREFIX}:proof:${digest(credential)}:${digest(nonce)}`; +} + +function validEd25519PublicKey(publicKey: string): boolean { + try { + return createPublicKey(publicKey).asymmetricKeyType === 'ed25519'; + } catch { + return false; + } +} + +export class RedisBridgePairingStore { + constructor( + private readonly redis: Redis, + private readonly pairingTtlSeconds = DEFAULT_PAIRING_TTL_SECONDS, + private readonly credentialTtlSeconds = DEFAULT_CREDENTIAL_TTL_SECONDS, + private readonly legacyScanClaimTtlMs = LEGACY_SCAN_CLAIM_TTL_MS, + private readonly rollbackEpoch = + process.env.CODEAPI_BRIDGE_PAIRING_ROLLBACK_EPOCH?.trim() ?? '', + ) {} + + async issue( + workerId: string, + binding?: BridgeWorkerBinding, + ): Promise { + // Pre-index binaries cannot remove a superseded code themselves. During + // the one pairing-TTL migration window, find and delete those records so + // rolling back cannot make a replaced code valid again. + await this.removeLegacyPairings(workerId); + const code = randomBytes(24).toString('base64url'); + const expiresAt = new Date( + Date.now() + this.pairingTtlSeconds * 1000, + ).toISOString(); + const generation = Number( + (await this.redis.get(workerPairingGenerationKey(workerId))) ?? '0', + ); + const pairing: StoredPairing = { workerId, expiresAt, generation, binding }; + const codeKey = pairingKey(code); + await this.redis.eval( + ISSUE_PAIRING_SCRIPT, + 2, + workerPairingIndexKey(workerId), + codeKey, + JSON.stringify(pairing), + String(this.pairingTtlSeconds), + ); + return { workerId, code, expiresAt }; + } + + async redeem(args: { + workerId: string; + code: string; + publicKey: string; + }): Promise { + const codeKey = pairingKey(args.code); + const raw = await this.redis.get(codeKey); + if (raw == null) { + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code is invalid or expired', + ); + } + const pairing = JSON.parse(raw) as StoredPairing; + if (pairing.workerId !== args.workerId) { + await this.redis.del(codeKey); + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code does not authorize this worker', + ); + } + if (!validEd25519PublicKey(args.publicKey)) { + throw new BridgePairingError( + 'PUBLIC_KEY_INVALID', + 'Worker public key must be an Ed25519 key', + ); + } + // Validate the supplied code before it can trigger a shared-keyspace scan. + // A rollback epoch means any generation-less code may have survived a + // legacy revoke, so clean the authenticated worker and reject that code. + if ( + pairing.generation == null && + this.rollbackEpoch.trim().length > 0 + ) { + await this.removeLegacyPairings(args.workerId); + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code is invalid or expired', + ); + } + + const credential = randomBytes(32).toString('base64url'); + const credentialDigest = digest(credential); + const expiresAt = new Date( + Date.now() + this.credentialTtlSeconds * 1000, + ).toISOString(); + const identityId = randomBytes(18).toString('base64url'); + const stored: StoredCredential = { + workerId: args.workerId, + identityId, + publicKey: args.publicKey, + expiresAt, + binding: pairing.binding, + }; + const accepted = await this.redis.eval( + REDEEM_PAIRING_SCRIPT, + 6, + codeKey, + workerPairingGenerationKey(pairing.workerId), + credentialDigestKey(credentialDigest), + workerIdentityKey(args.workerId), + workerPairingIndexKey(args.workerId), + workerStableIdentityKey(args.workerId), + raw, + pairing.generation == null ? '' : String(pairing.generation), + JSON.stringify(stored), + String(this.credentialTtlSeconds), + credentialDigest, + identityId, + ); + if (accepted !== 1) { + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code is invalid or expired', + ); + } + return { workerId: args.workerId, credential, expiresAt }; + } + + async authorize(args: { + workerId: string; + credential: string; + method: string; + path: string; + timestamp: string; + nonce: string; + body: string; + signature: string; + }): Promise<{ + workerId: string; + credentialId: string; + activeCredentialId: string; + identityId?: string; + pairingGeneration: number; + binding?: BridgeWorkerBinding; + }> { + const proofTime = Date.parse(args.timestamp); + if ( + !Number.isFinite(proofTime) || + Math.abs(Date.now() - proofTime) > PROOF_CLOCK_SKEW_MS + ) { + throw new BridgePairingError( + 'PROOF_INVALID', + 'Worker request proof is outside the accepted clock window', + ); + } + const credentialDigest = digest(args.credential); + const [raw, activeDigest, pairingGeneration] = await this.redis.mget( + credentialDigestKey(credentialDigest), + workerIdentityKey(args.workerId), + workerPairingGenerationKey(args.workerId), + ); + if (raw == null || activeDigest == null) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + const stored = JSON.parse(raw) as StoredCredential; + if (activeDigest !== credentialDigest) { + const activeRaw = await this.redis.get( + credentialDigestKey(activeDigest), + ); + const active = activeRaw == null + ? undefined + : JSON.parse(activeRaw) as StoredCredential; + if ( + stored.identityId == null || + active?.identityId == null || + stored.identityId !== active.identityId + ) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + } + if (stored.workerId !== args.workerId) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential does not authorize this worker', + ); + } + if (!verifyBridgeRequest(stored.publicKey, args, args.signature)) { + throw new BridgePairingError( + 'PROOF_INVALID', + 'Worker request proof is invalid', + ); + } + const accepted = await this.redis.set( + proofNonceKey(args.credential, args.nonce), + '1', + 'EX', + PROOF_NONCE_TTL_SECONDS, + 'NX', + ); + if (accepted !== 'OK') { + throw new BridgePairingError( + 'PROOF_REPLAYED', + 'Worker request proof has already been used', + ); + } + return { + workerId: stored.workerId, + credentialId: credentialDigest, + activeCredentialId: activeDigest, + ...(stored.identityId != null ? { identityId: stored.identityId } : {}), + pairingGeneration: Number(pairingGeneration ?? '0'), + ...(stored.binding ? { binding: stored.binding } : {}), + }; + } + + async revoke(workerId: string): Promise { + await this.removeLegacyPairings(workerId); + // Fence redemption and consume the currently indexed code atomically. An + // issue that linearized before this script is always removed; an issue + // that linearizes afterward installs a distinct generation and code. + await this.redis.eval( + REVOKE_PAIRING_SCRIPT, + 6, + workerPairingIndexKey(workerId), + workerPairingGenerationKey(workerId), + workerIdentityKey(workerId), + workerStableIdentityKey(workerId), + `${PREFIX}:worker:${encodeURIComponent(workerId)}`, + `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation`, + `${PREFIX}:credential:`, + `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:`, + ); + } + + private async removeLegacyPairings(workerId: string): Promise { + const deadlineKey = legacyPairingScanDeadlineKey(); + const now = Date.now(); + const migrationWindowMs = this.pairingTtlSeconds * 1000; + const proposedDeadline = now + migrationWindowMs; + let rawDeadline = await this.redis.get(deadlineKey); + if (rawDeadline == null) { + const initialized = await this.redis.set( + deadlineKey, + String(proposedDeadline), + 'NX', + ); + rawDeadline = initialized === 'OK' + ? String(proposedDeadline) + : await this.redis.get(deadlineKey); + } else if ((await this.redis.pttl(deadlineKey)) > 0) { + // Markers from the preceding build expired and reopened forever. Keep + // their original deadline, but make it durable so normal idle periods + // cannot start another migration window. + await this.redis.persist(deadlineKey); + } + + const indexedKey = await this.redis.get(workerPairingIndexKey(workerId)); + const indexedRaw = indexedKey == null ? null : await this.redis.get(indexedKey); + let rollbackDetected = false; + if (indexedRaw != null) { + try { + rollbackDetected = (JSON.parse(indexedRaw) as StoredPairing).generation == null; + } catch { + rollbackDetected = false; + } + } + + const deadline = Number(rawDeadline); + const stateKey = legacyPairingWorkerScanKey(workerId, this.rollbackEpoch); + const rollbackEpochDetected = this.rollbackEpoch.trim().length > 0; + while (true) { + const state = await this.redis.get(stateKey); + if (state === LEGACY_SCAN_COMPLETE && !rollbackDetected) return; + if (state === LEGACY_SCAN_PENDING) break; + if (state == null) { + if ( + !rollbackDetected && + !rollbackEpochDetected && + (!Number.isFinite(deadline) || Date.now() > deadline) + ) { + return; + } + const initialized = await this.redis.set( + stateKey, + LEGACY_SCAN_PENDING, + 'NX', + ); + if (initialized === 'OK') break; + continue; + } + // Predecessor builds stored an unqualified random token before scanning. + // It cannot prove whether that scan completed, so normalize it to a + // durable retry requirement instead of treating it as success. + const normalized = await this.redis.eval( + NORMALIZE_LEGACY_SCAN_STATE_SCRIPT, + 1, + stateKey, + state, + LEGACY_SCAN_PENDING, + ); + if (normalized === 1) break; + } + + const claimKey = `${stateKey}:claim`; + let scanClaim: { key: string; token: string } | undefined; + while (scanClaim == null) { + const state = await this.redis.get(stateKey); + if (state === LEGACY_SCAN_COMPLETE || state == null) return; + const token = `claim:${randomBytes(24).toString('base64url')}`; + const claimed = await this.redis.set( + claimKey, + token, + 'PX', + Math.max(1, this.legacyScanClaimTtlMs), + 'NX', + ); + if (claimed === 'OK') { + scanClaim = { key: claimKey, token }; + break; + } + await new Promise((resolve) => + setTimeout(resolve, LEGACY_SCAN_POLL_INTERVAL_MS), + ); + } + + let renewalError: unknown; + let renewal = Promise.resolve(); + let renewalInFlight = false; + const renewClaim = async (): Promise => { + const renewed = await this.redis.eval( + RENEW_LEGACY_SCAN_CLAIM_SCRIPT, + 1, + scanClaim.key, + scanClaim.token, + String(Math.max(1, this.legacyScanClaimTtlMs)), + ); + if (renewed !== 1) { + throw new Error('Legacy pairing cleanup claim was lost'); + } + }; + const renewalTimer = setInterval(() => { + if (renewalInFlight || renewalError != null) return; + renewalInFlight = true; + renewal = renewClaim() + .catch((error: unknown) => { + renewalError = error; + }) + .finally(() => { + renewalInFlight = false; + }); + }, Math.max(1, Math.floor(this.legacyScanClaimTtlMs / 3))); + renewalTimer.unref?.(); + + try { + let cursor = '0'; + do { + const [nextCursor, keys] = await this.redis.scan( + cursor, + 'MATCH', + `${PREFIX}:pairing:*`, + 'COUNT', + 100, + ); + if (renewalError != null) throw renewalError; + cursor = nextCursor; + if (keys.length === 0) continue; + const values = await this.redis.mget(...keys); + const matching = keys.filter((_key, index) => { + const raw = values[index]; + if (raw == null) return false; + try { + const pairing = JSON.parse(raw) as Partial; + return pairing.workerId === workerId && pairing.generation == null; + } catch { + return false; + } + }); + if (matching.length > 0) await this.redis.del(...matching); + } while (cursor !== '0'); + clearInterval(renewalTimer); + await renewal; + if (renewalError != null) throw renewalError; + await renewClaim(); + const completed = await this.redis.eval( + COMPLETE_LEGACY_SCAN_CLAIM_SCRIPT, + 2, + stateKey, + scanClaim.key, + scanClaim.token, + LEGACY_SCAN_COMPLETE, + String(deadline - Date.now()), + rollbackEpochDetected ? '1' : '0', + ); + if (completed !== 1) { + await this.removeLegacyPairings(workerId); + } + } catch (error) { + clearInterval(renewalTimer); + await renewal; + await this.redis.eval( + RELEASE_LEGACY_SCAN_CLAIM_SCRIPT, + 1, + scanClaim.key, + scanClaim.token, + ); + throw error; + } + } + + async rotate( + workerId: string, + expectedCredentialId?: string, + ): Promise { + const identityKey = workerIdentityKey(workerId); + const previousDigest = + expectedCredentialId ?? (await this.redis.get(identityKey)); + const previousRaw = + previousDigest == null + ? null + : await this.redis.get(credentialDigestKey(previousDigest)); + if (previousRaw == null || previousDigest == null) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + const previous = JSON.parse(previousRaw) as StoredCredential; + return await this.issueCredential( + workerId, + previous.publicKey, + previousDigest, + previous.binding, + previous.identityId ?? null, + ); + } + + private async issueCredential( + workerId: string, + publicKey: string, + previousDigest?: string, + binding?: BridgeWorkerBinding, + identityId?: string | null, + ): Promise { + const credential = randomBytes(32).toString('base64url'); + const credentialDigest = digest(credential); + const expiresAt = new Date( + Date.now() + this.credentialTtlSeconds * 1000, + ).toISOString(); + const stableIdentityId = + identityId === undefined + ? randomBytes(18).toString('base64url') + : identityId ?? undefined; + const stored: StoredCredential = { + workerId, + ...(stableIdentityId != null ? { identityId: stableIdentityId } : {}), + publicKey, + expiresAt, + binding, + }; + if (previousDigest !== undefined) { + const rotated = await this.redis.eval( + ROTATE_CREDENTIAL_SCRIPT, + 4, + workerIdentityKey(workerId), + credentialDigestKey(previousDigest), + credentialDigestKey(credentialDigest), + workerStableIdentityKey(workerId), + previousDigest, + credentialDigest, + JSON.stringify(stored), + String(this.credentialTtlSeconds), + stableIdentityId ?? '', + ); + if (rotated !== 1) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + } else { + const transaction = this.redis.multi(); + transaction.set( + credentialDigestKey(credentialDigest), + JSON.stringify(stored), + 'EX', + this.credentialTtlSeconds, + ); + transaction.set( + workerIdentityKey(workerId), + credentialDigest, + 'EX', + this.credentialTtlSeconds, + ); + if (stableIdentityId != null) { + transaction.set( + workerStableIdentityKey(workerId), + stableIdentityId, + 'EX', + this.credentialTtlSeconds, + ); + } else { + transaction.del(workerStableIdentityKey(workerId)); + } + await transaction.exec(); + } + return { workerId, credential, expiresAt }; + } +} diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts new file mode 100644 index 00000000..2b3baa10 --- /dev/null +++ b/service/src/bridge/router.test.ts @@ -0,0 +1,414 @@ +import { createServer, type Server } from 'http'; + +import { afterEach, describe, expect, test } from 'bun:test'; +import express, { json } from 'express'; +import RedisMock from 'ioredis-mock'; + +import type Redis from 'ioredis'; + +import { + createBridgeIdentity, + signBridgeRequest, +} from '../../../packages/code/src/identity'; +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { RedisBridgePairingStore } from './pairing'; +import { createBridgeRouter } from './router'; +import { RedisBridgeStore } from './store'; + +const redis = new RedisMock() as unknown as Redis; +let server: Server | undefined; + +afterEach(async () => { + server?.close(); + server = undefined; + await redis.flushall(); +}); + +describe('paired bridge HTTP API', () => { + test('rejects a malformed optional binding for a configured worker', async () => { + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store: new RedisBridgeStore(redis), + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + configuredWorkerId: 'vm-1', + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const response = await fetch( + `http://127.0.0.1:${address.port}/v1/bridge/pairings`, + { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + workerId: 'vm-1', + binding: { tenantId: 'tenant-1', principal: { type: 'user' } }, + }), + }, + ); + + expect(response.status).toBe(400); + }); + + test('requires and persists a trusted principal binding for dynamic workers', async () => { + const store = new RedisBridgeStore(redis); + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store, + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + allowDynamicWorkers: true, + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + const baseUrl = `http://127.0.0.1:${address.port}/v1/bridge`; + const unboundResponse = await fetch(`${baseUrl}/pairings`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'user-vm' }), + }); + expect(unboundResponse.status).toBe(400); + + const binding = { + tenantId: 'tenant-1', + principal: { type: 'user' as const, id: 'user-1' }, + }; + const pairingResponse = await fetch(`${baseUrl}/pairings`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'user-vm', binding }), + }); + const pairing = (await pairingResponse.json()) as { code: string }; + expect(pairingResponse.status).toBe(200); + + const identity = createBridgeIdentity(); + const redemptionResponse = await fetch(`${baseUrl}/pairings/redeem`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'user-vm', + code: pairing.code, + publicKey: identity.publicKey, + }), + }); + const issued = (await redemptionResponse.json()) as { credential: string }; + expect(redemptionResponse.status).toBe(200); + + const path = '/v1/bridge/workers/register'; + const body = JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'user-vm', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-2', + principal: { type: 'user', id: 'attacker-selected-user' }, + }, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path, + timestamp: new Date().toISOString(), + nonce: 'dynamic-registration-nonce', + body, + }; + const registrationResponse = await fetch( + `http://127.0.0.1:${address.port}${path}`, + { + method: 'POST', + headers: { + Authorization: `Bridge ${issued.credential}`, + 'Content-Type': 'application/json', + 'X-LibreChat-Code-Timestamp': proof.timestamp, + 'X-LibreChat-Code-Nonce': proof.nonce, + 'X-LibreChat-Code-Signature': signBridgeRequest( + identity.privateKey, + proof, + ), + }, + body, + }, + ); + expect(registrationResponse.status).toBe(200); + await expect( + store.dispatch({ + workerId: 'user-vm', + tenantId: binding.tenantId, + requireTenantBinding: true, + body: { language: 'bash' } as never, + headers: {}, + runtimeSessionId: 'stateful-session', + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + + await expect( + store.dispatch({ + workerId: 'user-vm', + tenantId: 'tenant-2', + requireTenantBinding: true, + body: { language: 'bash' } as never, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); + }); + + test('pairs a worker and accepts its proof-of-possession registration', async () => { + const app = express(); + const store = new RedisBridgeStore(redis); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store, + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + configuredWorkerId: 'vm-1', + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + const baseUrl = `http://127.0.0.1:${address.port}/v1/bridge`; + const pairingResponse = await fetch(`${baseUrl}/pairings`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'vm-1' }), + }); + const pairing = (await pairingResponse.json()) as { code: string }; + expect(pairingResponse.status).toBe(200); + + const identity = createBridgeIdentity(); + const redemptionResponse = await fetch(`${baseUrl}/pairings/redeem`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }), + }); + const issued = (await redemptionResponse.json()) as { + credential: string; + }; + expect(redemptionResponse.status).toBe(200); + + const path = '/v1/bridge/workers/register'; + const body = JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path, + timestamp: new Date().toISOString(), + nonce: 'http-registration-nonce', + body, + }; + const headers = { + Authorization: `Bridge ${issued.credential}`, + 'Content-Type': 'application/json', + 'X-LibreChat-Code-Timestamp': proof.timestamp, + 'X-LibreChat-Code-Nonce': proof.nonce, + 'X-LibreChat-Code-Signature': signBridgeRequest( + identity.privateKey, + proof, + ), + }; + const registrationUrl = `http://127.0.0.1:${address.port}${path}`; + const registrationResponse = await fetch(registrationUrl, { + method: 'POST', + headers, + body, + }); + + expect(registrationResponse.status).toBe(200); + await expect(registrationResponse.json()).resolves.toMatchObject({ + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + }); + + const crossDeploymentRevoke = await fetch( + `${baseUrl}/workers/another-deployments-worker/revoke`, + { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: '{}', + }, + ); + expect(crossDeploymentRevoke.status).toBe(400); + + const replayResponse = await fetch(registrationUrl, { + method: 'POST', + headers, + body, + }); + expect(replayResponse.status).toBe(401); + await expect(replayResponse.json()).resolves.toMatchObject({ + code: 'PROOF_REPLAYED', + }); + + const revokeResponse = await fetch(`${baseUrl}/workers/vm-1/revoke`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: '{}', + }); + expect(revokeResponse.status).toBe(200); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_FENCED' }); + }); + + test('forwards pairing store failures to Express error middleware', async () => { + const app = express(); + const pairings = new RedisBridgePairingStore(redis); + pairings.issue = async () => { + throw new Error('pairing store unavailable'); + }; + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store: new RedisBridgeStore(redis), + pairings, + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + configuredWorkerId: 'vm-1', + }), + ); + app.use( + ( + error: Error, + _req: express.Request, + res: express.Response, + _next: express.NextFunction, + ) => { + res.status(503).json({ error: error.message }); + }, + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const response = await fetch( + `http://127.0.0.1:${address.port}/v1/bridge/pairings`, + { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'vm-1' }), + }, + ); + + expect(response.status).toBe(503); + await expect(response.json()).resolves.toEqual({ + error: 'pairing store unavailable', + }); + }); + + test('does not treat a missing configured worker ID as a wildcard', async () => { + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store: new RedisBridgeStore(redis), + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const response = await fetch( + `http://127.0.0.1:${address.port}/v1/bridge/pairings`, + { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'vm-1' }), + }, + ); + + expect(response.status).toBe(400); + }); +}); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts new file mode 100644 index 00000000..e41336d4 --- /dev/null +++ b/service/src/bridge/router.ts @@ -0,0 +1,659 @@ +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 { BridgePrincipalType, BridgeWorkerBinding } from './pairing'; +import type { CodeBridgeAssignment, CodeBridgeSettlement } from './store'; + +import { + BRIDGE_PROTOCOL_VERSION, + isValidBridgeWorkerCapabilities, + isValidBridgeWorkerId, +} from '../../../packages/code/src/protocol'; +import { BridgePairingError, RedisBridgePairingStore } from './pairing'; +import { BridgeStoreError, RedisBridgeStore } from './store'; + +const INCARNATION_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; +const MAX_LEASE_WAIT_MS = 30_000; +const BRIDGE_BINDING_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; +const PRINCIPAL_TYPES = new Set([ + 'deployment', + 'tenant', + 'user', + 'role', + 'group', +]); + +export type BridgeAuthMode = 'static' | 'paired'; + +export interface BridgeRouterOptions { + store: RedisBridgeStore; + pairings: RedisBridgePairingStore; + authMode: BridgeAuthMode; + adminToken: string; + configuredWorkerId?: string; + allowDynamicWorkers?: boolean; +} + +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 validWorkerId(value: string): boolean { + return isValidBridgeWorkerId(value); +} + +function validIncarnationId(value: unknown): value is string { + return typeof value === 'string' && INCARNATION_ID_PATTERN.test(value); +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function parseBinding(value: unknown): BridgeWorkerBinding | undefined { + if (!isRecord(value) || !isRecord(value.principal)) return undefined; + const { tenantId, principal } = value; + if ( + typeof tenantId !== 'string' || + !BRIDGE_BINDING_ID_PATTERN.test(tenantId) || + typeof principal.type !== 'string' || + !PRINCIPAL_TYPES.has(principal.type as BridgePrincipalType) || + typeof principal.id !== 'string' || + !BRIDGE_BINDING_ID_PATTERN.test(principal.id) + ) { + return undefined; + } + return { + tenantId, + principal: { + type: principal.type as BridgePrincipalType, + id: principal.id, + }, + }; +} + +function asyncRoute( + handler: (req: Request, res: Response) => Promise, +): RequestHandler { + return (req, res, next) => { + void handler(req, res).catch(next); + }; +} + +function sendStoreError(error: BridgeStoreError, res: Response): void { + const status = + error.code === 'ASSIGNMENT_NOT_FOUND' + ? 404 + : error.code === 'WORKER_UNAUTHORIZED' + ? 403 + : error.code === 'WORKER_BUSY' + ? 503 + : 409; + res.status(status).json({ error: error.message, code: error.code }); +} + +function isSettlement(value: unknown): value is CodeBridgeSettlement { + if (!isRecord(value)) return false; + if ( + value.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof value.generation !== 'number' || + !Number.isSafeInteger(value.generation) || + value.generation < 1 || + typeof value.leaseToken !== 'string' || + value.leaseToken.length < 32 || + !validIncarnationId(value.incarnationId) + ) { + return false; + } + if (value.status === 'rejected') { + return typeof value.error === 'string' && value.error.length <= 4096; + } + return ( + value.status === 'fulfilled' && + typeof value.result === 'object' && + value.result !== null + ); +} + +export function createBridgeRouter(options: BridgeRouterOptions): Router { + const router = Router(); + + const configuredWorker = (workerId: string): boolean => + options.allowDynamicWorkers === true || + (options.configuredWorkerId != null && + options.configuredWorkerId !== '' && + workerId === options.configuredWorkerId); + + const bearerToken = (req: Request): string => + req + .header('Authorization') + ?.match(/^Bearer\s+(.+)$/i)?.[1] + ?.trim() ?? ''; + + const adminAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { + if (!options.adminToken) { + res.status(503).json({ error: 'Code bridge is not configured' }); + return; + } + const token = bearerToken(req); + if (!token || !sameToken(token, options.adminToken)) { + res.status(401).json({ error: 'Invalid code bridge administrator token' }); + return; + } + next(); + }; + + const staticWorkerAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { + const token = bearerToken(req); + if (!token || !sameToken(token, options.adminToken)) { + res.status(401).json({ error: 'Invalid code bridge worker token' }); + return; + } + next(); + }; + + const pairedWorkerAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { + const workerId = + req.params.workerId || + (isRecord(req.body) && typeof req.body.workerId === 'string' + ? req.body.workerId + : ''); + const credential = + req + .header('Authorization') + ?.match(/^Bridge\s+(.+)$/i)?.[1] + ?.trim() ?? ''; + const timestamp = req.header('X-LibreChat-Code-Timestamp') ?? ''; + const nonce = req.header('X-LibreChat-Code-Nonce') ?? ''; + const signature = req.header('X-LibreChat-Code-Signature') ?? ''; + if ( + !validWorkerId(workerId) || + !credential || + !timestamp || + !nonce || + !signature + ) { + res.status(401).json({ error: 'Invalid paired worker authorization' }); + return; + } + void options.pairings + .authorize({ + workerId, + credential, + method: req.method, + path: req.originalUrl.split('?')[0], + timestamp, + nonce, + body: JSON.stringify(req.body ?? {}), + signature, + }) + .then((authorization) => { + res.locals.bridgeWorkerAuthorization = authorization; + next(); + }) + .catch((error: unknown) => { + if (error instanceof BridgePairingError) { + res.status(401).json({ error: error.message, code: error.code }); + return; + } + next(error); + }); + }; + + const workerAuth = + options.authMode === 'paired' ? pairedWorkerAuth : staticWorkerAuth; + + router.post('/pairings', adminAuth, asyncRoute(async (req, res) => { + if (options.authMode !== 'paired') { + res.status(409).json({ error: 'Paired worker authentication is disabled' }); + return; + } + const workerId = isRecord(req.body) ? req.body.workerId : undefined; + if ( + typeof workerId !== 'string' || + !validWorkerId(workerId) || + !configuredWorker(workerId) + ) { + res.status(400).json({ error: 'Invalid bridge worker ID' }); + return; + } + const hasBinding = isRecord(req.body) && + Object.prototype.hasOwnProperty.call(req.body, 'binding'); + const binding = isRecord(req.body) ? parseBinding(req.body.binding) : undefined; + if (hasBinding && binding == null) { + res.status(400).json({ error: 'Invalid bridge worker principal binding' }); + return; + } + if (options.allowDynamicWorkers === true && binding == null) { + res.status(400).json({ + error: 'Dynamic bridge workers require a valid principal binding', + }); + return; + } + const pairing = await options.pairings.issue(workerId, binding); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); + })); + + router.post('/pairings/redeem', asyncRoute(async (req, res) => { + if (options.authMode !== 'paired') { + res.status(409).json({ error: 'Paired worker authentication is disabled' }); + return; + } + const redemption = req.body as unknown; + if ( + !isRecord(redemption) || + redemption.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof redemption.workerId !== 'string' || + !validWorkerId(redemption.workerId) || + !configuredWorker(redemption.workerId) || + typeof redemption.code !== 'string' || + redemption.code.length < 16 || + typeof redemption.publicKey !== 'string' || + redemption.publicKey.length > 4096 + ) { + res.status(400).json({ error: 'Invalid bridge pairing redemption' }); + return; + } + try { + const credential = await options.pairings.redeem({ + workerId: redemption.workerId, + code: redemption.code, + publicKey: redemption.publicKey, + }); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); + } catch (error) { + if (error instanceof BridgePairingError) { + const status = error.code === 'PUBLIC_KEY_INVALID' ? 400 : 401; + res.status(status).json({ error: error.message, code: error.code }); + return; + } + throw error; + } + })); + + router.post( + '/workers/:workerId/revoke', + adminAuth, + asyncRoute(async (req, res) => { + if ( + !validWorkerId(req.params.workerId) || + !configuredWorker(req.params.workerId) + ) { + res.status(400).json({ error: 'Invalid bridge worker ID' }); + return; + } + await options.pairings.revoke(req.params.workerId); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, revoked: true }); + }), + ); + + +router.post( + '/workers/register', + workerAuth, + 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) || + !isValidBridgeWorkerCapabilities(registration.capabilities) + ) { + res.status(400).json({ error: 'Invalid bridge worker registration' }); + return; + } + if ( + !configuredWorker(registration.workerId) + ) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + const authorization = options.authMode === 'paired' + ? ( + res.locals.bridgeWorkerAuthorization as { + identityId: string; + pairingGeneration: number; + credentialId: string; + activeCredentialId: string; + binding?: BridgeWorkerBinding; + } + ) + : undefined; + const trustedRegistration: BridgeWorkerRegistration & { + binding?: BridgeWorkerBinding; + credentialId?: string; + identityId?: string; + } = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + capabilities: registration.capabilities, + ...(authorization?.credentialId != null + ? { credentialId: authorization.credentialId } + : {}), + ...(authorization?.identityId != null + ? { identityId: authorization.identityId } + : {}), + ...(authorization?.binding != null + ? { binding: authorization.binding } + : {}), + }; + try { + await options.store.register( + trustedRegistration, + authorization, + ); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }), +); + +router.post( + '/workers/:workerId/workspaces/reset', + workerAuth, + 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 (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { + const resetController = new AbortController(); + const abortReset = (): void => resetController.abort(); + req.once('aborted', abortReset); + res.once('close', abortReset); + try { + await options.store.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); + return; + } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/lease', + workerAuth, + 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); + 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 (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { + const leaseController = new AbortController(); + const abortLease = (): void => leaseController.abort(); + req.once('aborted', abortLease); + res.once('close', abortLease); + let assignment: CodeBridgeAssignment | undefined; + try { + assignment = await options.store.lease( + workerId, + body.incarnationId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + leaseController.signal, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, + ); + if (leaseController.signal.aborted) { + if (assignment != null) await options.store.returnLease(assignment); + return; + } + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), + assignment, + }); + } finally { + req.off('aborted', abortLease); + res.off('close', abortLease); + } + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/assignments/:assignmentId/ack', + workerAuth, + 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 { + const acknowledgementController = new AbortController(); + const abortAcknowledgement = (): void => + acknowledgementController.abort(); + req.once('aborted', abortAcknowledgement); + res.once('close', abortAcknowledgement); + try { + await options.store.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); + return; + } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/assignments/:assignmentId/settle', + workerAuth, + asyncRoute(async (req, res) => { + const settlement = req.body as unknown; + if (!isSettlement(settlement)) { + res.status(400).json({ error: 'Invalid bridge settlement' }); + return; + } + try { + const settlementController = new AbortController(); + const abortSettlement = (): void => settlementController.abort(); + req.once('aborted', abortSettlement); + res.once('close', abortSettlement); + try { + await options.store.settle( + req.params.workerId, + req.params.assignmentId, + settlement, + settlementController.signal, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, + ); + 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); + return; + } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/assignments/:assignmentId/cancellation', + workerAuth, + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) + ) { + res.status(400).json({ error: 'Invalid bridge cancellation request' }); + return; + } + const cancellationController = new AbortController(); + const abortCancellation = (): void => cancellationController.abort(); + req.once('aborted', abortCancellation); + res.once('close', abortCancellation); + try { + const cancelled = await options.store.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); + } + }), +); + + router.post( + '/workers/:workerId/credentials/refresh', + workerAuth, + asyncRoute(async (req, res) => { + try { + const credential = await options.pairings.rotate( + req.params.workerId, + ( + res.locals.bridgeWorkerAuthorization as + | { credentialId: string } + | undefined + )?.credentialId, + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); + } catch (error) { + if (error instanceof BridgePairingError) { + res.status(401).json({ error: error.message, code: error.code }); + return; + } + throw error; + } + }), + ); + + + return router; +} diff --git a/service/src/bridge/selection.test.ts b/service/src/bridge/selection.test.ts new file mode 100644 index 00000000..b6309702 --- /dev/null +++ b/service/src/bridge/selection.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, test } from 'bun:test'; + +import { + BridgeWorkerSelectionError, + resolveBridgeWorkerSelection, +} from './selection'; + +describe('bridge worker request selection', () => { + test('uses the configured compatibility worker when no dynamic worker is requested', () => { + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + }), + ).toEqual({ workerId: 'deployment-worker', explicit: false }); + }); + + test('selects only the worker authenticated by the LibreChat JWT', () => { + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'code-user_1', + trustedWorkerId: 'code-user_1', + }), + ).toEqual({ workerId: 'code-user_1', explicit: true }); + + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + trustedWorkerId: 'code-user_1', + }), + ).toEqual({ workerId: 'code-user_1', explicit: true }); + }); + + test('rejects a caller-controlled worker header without a matching trusted claim', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'victim-worker', + }), + ).toThrow('Code bridge worker selection is not authenticated'); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'victim-worker', + trustedWorkerId: 'caller-worker', + }), + ).toThrow('Code bridge worker selection does not match the authenticated claim'); + }); + + test('rejects dynamic routing on the wrong backend or when it is disabled', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'http', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: 'code-user-1', + trustedWorkerId: 'code-user-1', + }), + ).toThrow(BridgeWorkerSelectionError); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: false, + requestedWorkerId: 'code-user-1', + trustedWorkerId: 'code-user-1', + }), + ).toThrow('Dynamic code bridge workers are disabled'); + }); + + test('rejects malformed worker IDs before they cross the queue boundary', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: '../worker', + trustedWorkerId: '../worker', + }), + ).toThrow('Invalid code bridge worker ID'); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: 'victim:assignments', + trustedWorkerId: 'victim:assignments', + }), + ).toThrow('Invalid code bridge worker ID'); + }); +}); diff --git a/service/src/bridge/selection.ts b/service/src/bridge/selection.ts new file mode 100644 index 00000000..0959279f --- /dev/null +++ b/service/src/bridge/selection.ts @@ -0,0 +1,72 @@ +export const CODEAPI_BRIDGE_WORKER_HEADER = 'X-LibreChat-Code-Worker-ID'; +export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/; + +export class BridgeWorkerSelectionError extends Error { + constructor( + message: string, + public readonly status: 400 | 403 | 503, + ) { + super(message); + this.name = 'BridgeWorkerSelectionError'; + } +} + +export function resolveBridgeWorkerSelection(args: { + backend: SandboxBackendName; + configuredWorkerId: string; + dynamicWorkers: boolean; + requestedWorkerId?: string; + trustedWorkerId?: string; +}): { workerId: string; explicit: boolean } | undefined { + const requestedWorkerId = args.requestedWorkerId?.trim(); + const trustedWorkerId = args.trustedWorkerId?.trim(); + const hasRequestedWorker = requestedWorkerId != null && requestedWorkerId.length > 0; + const hasTrustedWorker = trustedWorkerId != null && trustedWorkerId.length > 0; + if (hasRequestedWorker || hasTrustedWorker) { + if (args.backend !== 'remote-bridge') { + throw new BridgeWorkerSelectionError( + 'Code bridge worker routing requires the remote-bridge backend', + 400, + ); + } + if (hasRequestedWorker && !hasTrustedWorker) { + throw new BridgeWorkerSelectionError( + 'Code bridge worker selection is not authenticated', + 403, + ); + } + if ( + hasRequestedWorker && + hasTrustedWorker && + requestedWorkerId !== trustedWorkerId + ) { + throw new BridgeWorkerSelectionError( + 'Code bridge worker selection does not match the authenticated claim', + 403, + ); + } + const selectedWorkerId = trustedWorkerId as string; + if (!BRIDGE_WORKER_ID_PATTERN.test(selectedWorkerId)) { + throw new BridgeWorkerSelectionError('Invalid code bridge worker ID', 400); + } + if (!args.dynamicWorkers && selectedWorkerId !== args.configuredWorkerId) { + throw new BridgeWorkerSelectionError('Dynamic code bridge workers are disabled', 403); + } + return { + workerId: selectedWorkerId, + explicit: true, + }; + } + + if (args.backend !== 'remote-bridge') return undefined; + const configuredWorkerId = args.configuredWorkerId.trim(); + if (configuredWorkerId.length === 0) { + throw new BridgeWorkerSelectionError('No code bridge worker was selected', 503); + } + if (!BRIDGE_WORKER_ID_PATTERN.test(configuredWorkerId)) { + throw new BridgeWorkerSelectionError('Invalid configured code bridge worker ID', 503); + } + return { workerId: configuredWorkerId, explicit: false }; +} + +type SandboxBackendName = 'http' | 'lambda-microvm' | 'remote-bridge'; diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts new file mode 100644 index 00000000..23b04067 --- /dev/null +++ b/service/src/bridge/store.test.ts @@ -0,0 +1,2077 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { getEventListeners } from 'node:events'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import type * as t from '../types'; +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { RedisBridgeStore } from './store'; + +const redis = new RedisMock() as unknown as Redis; +const store = new RedisBridgeStore(redis); +const incarnationId = 'incarnation-00000001'; +const redisEval = redis.eval.bind(redis); +const redisDel = redis.del.bind(redis); +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(); +}); + +describe('RedisBridgeStore', () => { + test('rejects a registration whose authenticated identity was replaced', async () => { + await redis.set( + 'codeapi:bridge:v1:identity:fenced-worker', + 'replacement-credential-digest', + ); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'fenced-worker', + incarnationId, + credentialId: 'stale-credential-digest', + identityId: 'stale-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }, 'stale-credential-digest'), + ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); + + await expect( + redis.get('codeapi:bridge:v1:worker:fenced-worker'), + ).resolves.toBeNull(); + }); + + test('accepts registration after a same-identity credential rotation', async () => { + await redis.set( + 'codeapi:bridge:v1:identity:rotating-registration-worker', + 'new-active-credential-digest', + ); + await redis.set( + 'codeapi:bridge:v1:stable-identity:rotating-registration-worker', + 'stable-worker-identity', + ); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rotating-registration-worker', + incarnationId, + credentialId: 'old-authenticated-credential-digest', + identityId: 'stable-worker-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }, 'old-authenticated-credential-digest'), + ).resolves.toBeUndefined(); + }); + + test('rejects a dynamic worker lease outside its bound tenant', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'tenant-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-1', + principal: { type: 'user', id: 'user-1' }, + }, + }); + + await expect( + store.dispatch({ + workerId: 'tenant-worker', + tenantId: 'tenant-2', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); + }); + + test('does not lease an assignment to a newly rebound worker identity', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rebound-worker', + incarnationId, + identityId: 'tenant-a-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'rebound-worker', + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect( + store.lease( + 'rebound-worker', + incarnationId, + 1_000, + undefined, + 'tenant-b-identity', + ), + ).resolves.toBeUndefined(); + await expect( + store.lease( + 'rebound-worker', + incarnationId, + 1_000, + undefined, + 'tenant-a-identity', + ), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('a stale identity poll cannot consume work queued for the replacement identity', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'replacement-worker', + incarnationId, + identityId: 'replacement-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'replacement-worker', + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect( + store.lease( + 'replacement-worker', + incarnationId, + 100, + undefined, + 'stale-identity', + ), + ).resolves.toBeUndefined(); + await expect( + store.lease( + 'replacement-worker', + incarnationId, + 1_000, + undefined, + 'replacement-identity', + ), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('a stale incarnation poll cannot consume replacement incarnation work', async () => { + const replacementIncarnationId = 'incarnation-00000002'; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'restarted-worker', + incarnationId: replacementIncarnationId, + identityId: 'stable-restarted-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + 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( + store.lease( + 'restarted-worker', + incarnationId, + 100, + undefined, + 'stable-restarted-identity', + ), + ).resolves.toBeUndefined(); + await expect( + store.lease( + 'restarted-worker', + replacementIncarnationId, + 1_000, + undefined, + 'stable-restarted-identity', + ), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('leases queued work after credential refresh preserves the paired identity', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rotating-worker', + incarnationId, + identityId: 'stable-paired-identity', + credentialId: 'credential-before-refresh', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'rotating-worker', + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + const assignment = await store.lease( + 'rotating-worker', + incarnationId, + 1_000, + undefined, + 'stable-paired-identity', + ); + + expect(assignment).toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('delivers and settles one fenced stateful assignment', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: { 'X-Execution-Manifest': 'signed' }, + runtimeSessionId: 'rt-user-1', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('vm-1', incarnationId, 1_000); + expect(assignment).toBeDefined(); + expect(assignment?.runtimeSessionId).toBe('rt-user-1'); + expect(assignment?.remainingMs).toBeGreaterThan(0); + expect(assignment?.remainingMs).toBeLessThanOrEqual(5_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).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'run-1' }, + }); + }); + + 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('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 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 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('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('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('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, + 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, + 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']; + + 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( + 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('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('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('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, + 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, + 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.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.eval = redisEval as Redis['eval']; + + 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('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, + 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, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'vm-1', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('vm-1', incarnationId, 1_000); + + await expect( + store.settle('vm-1', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: 'stale-token-that-is-long-enough-to-pass-validation', + incarnationId, + status: 'rejected', + error: 'unused', + }), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('fences a replaced worker incarnation', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_FENCED' }); + }); + + test('a stale incarnation poll cannot consume replacement work', async () => { + const replacementIncarnationId = 'incarnation-00000002'; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'restarted-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const stalePoll = store.lease('restarted-worker', incarnationId, 100); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'restarted-worker', + incarnationId: replacementIncarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'restarted-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect(stalePoll).resolves.toBeUndefined(); + await expect( + store.lease('restarted-worker', replacementIncarnationId, 1_000), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('dispatch retries atomically against a replacement incarnation', async () => { + const workerId = 'racing-worker'; + const replacementIncarnationId = 'incarnation-00000002'; + const capabilities = { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [] as string[], + }; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities, + }); + const originalEval = redis.eval.bind(redis); + let replaced = false; + redis.eval = (async (...args: Parameters) => { + if (!replaced && String(args[0]).includes("redis.call('RPUSH'")) { + replaced = true; + const replacement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId: replacementIncarnationId, + capabilities, + }; + await redis.set( + `codeapi:bridge:v1:worker:${workerId}`, + JSON.stringify(replacement), + 'EX', + 60, + ); + await redis.set( + `codeapi:bridge:v1:worker:${workerId}:incarnation`, + replacementIncarnationId, + 'EX', + 60, + ); + } + return originalEval(...args); + }) as Redis['eval']; + const controller = new AbortController(); + const completion = store.dispatch({ + workerId, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + const assignment = await store.lease( + workerId, + replacementIncarnationId, + 1_000, + ); + expect(assignment?.incarnationId).toBe(replacementIncarnationId); + redis.eval = originalEval as Redis['eval']; + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('defers worker replacement while an assignment is active', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'busy-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'busy-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease('busy-worker', incarnationId, 1_000); + expect(assignment).toBeDefined(); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'busy-worker', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'busy-worker', + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).resolves.toBeUndefined(); + }); + + test('recovers only the assignment owner after registration expiry', async () => { + const workerId = 'expired-registration-worker'; + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease(workerId, incarnationId, 1_000); + expect(assignment).toBeDefined(); + await redis.del( + `codeapi:bridge:v1:worker:${workerId}`, + `codeapi:bridge:v1:worker:${workerId}:incarnation`, + ); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId: 'incarnation-00000002', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }), + ).resolves.toBeUndefined(); + + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('removes abort listeners after each settlement poll delay', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'listener-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'listener-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 350, + signal: controller.signal, + }); + + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + expect(getEventListeners(controller.signal, 'abort')).toHaveLength(0); + }); + + test('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('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, + workerId: 'long-running-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'long-running-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 15 * 60_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + const [assignmentKey] = await redis.keys('codeapi:bridge:v1:assignment:*'); + + expect(await redis.ttl(assignmentKey)).toBeGreaterThan(10 * 60); + expect( + await redis.pttl('codeapi:bridge:v1:worker:long-running-worker:lock'), + ).toBeGreaterThan(10 * 60_000); + controller.abort(); + await expect(completion).rejects.toMatchObject({ + code: 'ASSIGNMENT_EXPIRED', + }); + }); + + test('observes a settlement accepted during the final poll delay', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'deadline-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const deadlineAtMs = Date.now() + 500; + const completion = store.dispatch({ + workerId: 'deadline-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-deadline', + deadlineAtMs, + signal: controller.signal, + }); + const assignment = await store.lease( + 'deadline-worker', + incarnationId, + 1_000, + ); + expect(assignment).toBeDefined(); + await new Promise((resolve) => + setTimeout(resolve, Math.max(0, deadlineAtMs - Date.now() - 30)), + ); + await store.settle('deadline-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + language: 'bash', + version: '5.2.0', + session_id: 'run-deadline', + files: [], + }, + }); + + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'run-deadline' }, + }); + }); + + test('preserves a committed result across transient cleanup failures', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'cleanup-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: [], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'cleanup-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + runtimeSessionId: 'rt-cleanup', + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'cleanup-worker', + incarnationId, + 1_000, + ); + let cleanupAttempts = 0; + 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, + 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.eval = redisEval as Redis['eval']; + }); + + 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, + ); + const settlement = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + 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', + ); + 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); + await expect( + store.settle( + 'commit-worker', + assignment?.assignmentId ?? '', + settlement, + ), + ).resolves.toBeUndefined(); + 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, + 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(); + 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', + ); + 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({ + 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', + 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' }); + + 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 () => { + 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('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 store.acknowledgeLease( + 'late-rejection-worker', + incarnationId, + assignment?.assignmentId ?? '', + assignment?.generation ?? 0, + assignment?.leaseToken ?? '', + ); + 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('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; + 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, + 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 ?? '', + ); + 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, + 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: 'WORKSPACE_QUARANTINED' }); + 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' }); + }); + + 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 new file mode 100644 index 00000000..3a43298b --- /dev/null +++ b/service/src/bridge/store.ts @@ -0,0 +1,1424 @@ +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'; +import type { BridgeWorkerBinding } from './pairing'; + +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< + t.ExecuteResponse & { + session_id: string; + files?: t.FileRefs; + run?: t.ExecuteResponse['run']; + } +>; + +export class BridgeStoreError extends Error { + constructor( + public readonly code: + | 'WORKER_OFFLINE' + | 'WORKER_UNAUTHORIZED' + | 'WORKER_BUSY' + | 'ASSIGNMENT_EXPIRED' + | 'ASSIGNMENT_FENCED' + | 'ASSIGNMENT_NOT_FOUND' + | 'WORKER_FENCED' + | 'WORKER_QUARANTINED' + | 'WORKSPACE_QUARANTINED' + | 'WORKER_MISMATCH', + message: string, + ) { + super(message); + this.name = 'BridgeStoreError'; + } +} + +interface StoredAssignment extends CodeBridgeAssignment { + leaseTokenHash: string; + workerIdentityId?: string; +} + +export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { + binding?: BridgeWorkerBinding; + credentialId?: string; + identityId?: string; +} + +function workerKey(workerId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}`; +} + +function workerStableIdentityKey(workerId: string): string { + return `${PREFIX}:stable-identity:${workerId}`; +} + +function workerIncarnationKey(workerId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation`; +} + +function incarnationFenceKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:fenced`; +} + +function quarantineKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:quarantined`; +} + +function workspaceQuarantineKey( + workerId: string, + runtimeSessionId: string, +): string { + const sessionHash = createHash('sha256') + .update(runtimeSessionId) + .digest('hex'); + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:workspace:${sessionHash}:quarantined`; +} + +function queueKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:assignments`; +} + +function leaseClaimKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-claim`; +} + +function leaseAckKey(workerId: string, incarnationId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-ack`; +} + +function generationKey(workerId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:generation`; +} + +function lockKey(workerId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:lock`; +} + +function lockIncarnationKey(workerId: string): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:lock:incarnation`; +} + +function assignmentKey(assignmentId: string): string { + return `${PREFIX}:assignment:${assignmentId}`; +} + +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`; +} + +function tokenHash(token: string): string { + return createHash('sha256').update(token).digest('hex'); +} + +function assignmentTtlSeconds(deadlineAtMs: number): number { + return Math.max(1, Math.ceil((deadlineAtMs - Date.now()) / 1000) + 30); +} + +async function delay(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted === true) return; + await new Promise((resolve) => { + const onAbort = (): void => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} + +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, + ) {} + + 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; + } + } + + private async leaseCommand( + command: Promise, + signal: AbortSignal | undefined, + label: string, + ): Promise { + return await boundedCommand( + command, + this.redisCommandTimeoutMs, + label, + signal, + ); + } + + async register( + registration: RegisteredBridgeWorker, + authorization?: string | { + identityId?: string; + pairingGeneration?: number; + activeCredentialId?: string; + }, + ): Promise { + const authorizationObject = + typeof authorization === 'object' ? authorization : undefined; + const expectedActiveCredentialId = + typeof authorization === 'string' + ? authorization + : authorizationObject?.activeCredentialId; + const script = [ + 'if ARGV[5] ~= "" then', + ' local pairingGeneration = redis.call(\'GET\', KEYS[7]) or "0"', + ' if pairingGeneration ~= ARGV[5] then return -5 end', + ' if ARGV[6] ~= "" then', + ' if redis.call(\'GET\', KEYS[8]) ~= ARGV[6] then return -5 end', + ' elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -5', + ' end', + 'end', + 'if ARGV[8] ~= "" then', + ' local stableIdentity = redis.call(\'GET\', KEYS[8])', + ' if stableIdentity and stableIdentity ~= ARGV[8] then return -4 end', + ' if not stableIdentity then', + ' if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4 end', + ' redis.call(\'SET\', KEYS[8], ARGV[8], "EX", ARGV[3])', + ' end', + 'elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4', + 'end', + 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', + 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', + 'local current = redis.call(\'GET\', KEYS[4])', + 'if not current and redis.call(\'EXISTS\', KEYS[5]) == 1 then', + ' local owner = redis.call(\'GET\', KEYS[6])', + ' if owner ~= ARGV[1] then return -3 end', + 'end', + 'if current then', + ' if current ~= ARGV[1] then', + ' if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -3 end', + ' redis.call(\'SET\', ARGV[4] .. current .. \':fenced\', \"1\")', + ' end', + 'end', + 'redis.call(\'SET\', KEYS[1], ARGV[2], \"EX\", ARGV[3])', + 'redis.call(\'SET\', KEYS[4], ARGV[1], \"EX\", ARGV[3])', + 'return 1', + ].join('\n'); + const result = Number( + await boundedCommand( + this.redis.eval( + script, + 9, + workerKey(registration.workerId), + incarnationFenceKey(registration.workerId, registration.incarnationId), + quarantineKey(registration.workerId, registration.incarnationId), + workerIncarnationKey(registration.workerId), + lockKey(registration.workerId), + lockIncarnationKey(registration.workerId), + `${PREFIX}:pairing-generation:${registration.workerId}`, + `${PREFIX}:stable-identity:${registration.workerId}`, + `${PREFIX}:identity:${registration.workerId}`, + registration.incarnationId, + JSON.stringify(registration), + String(this.workerTtlSeconds), + `${PREFIX}:worker:${encodeURIComponent(registration.workerId)}:incarnation:`, + authorizationObject?.pairingGeneration == null + ? '' + : String(authorizationObject.pairingGeneration), + authorizationObject?.identityId ?? '', + expectedActiveCredentialId ?? '', + registration.identityId ?? '', + ), + this.redisCommandTimeoutMs, + 'Bridge worker registration', + ), + ); + if (result === -2) { + throw new BridgeStoreError( + 'WORKER_QUARANTINED', + 'Bridge worker incarnation is quarantined', + ); + } + if (result === -1) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Bridge worker incarnation was replaced', + ); + } + if (result === -3) { + throw new BridgeStoreError( + 'WORKER_BUSY', + 'Bridge worker cannot be replaced during an active assignment', + ); + } + if (result === -4) { + throw new BridgeStoreError( + 'WORKER_UNAUTHORIZED', + 'Bridge worker authorization was revoked before registration completed', + ); + } + if (result === -5) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Bridge worker authorization was revoked before registration completed', + ); + } + } + + async dispatch(args: { + workerId: string; + tenantId?: string; + requireTenantBinding?: boolean; + body: t.PayloadBody; + headers: Record; + runtimeSessionId?: string; + deadlineAtMs: number; + signal: AbortSignal; + finalize?: ( + settlement: CodeBridgeSettlement, + ) => Promise; + }): Promise { + this.assertDispatchActive(args.signal, args.deadlineAtMs); + let registration = await this.dispatchCommand( + () => this.registration(args.workerId), + args, + 'Bridge worker registration read', + ); + if (registration == null) { + throw new BridgeStoreError( + 'WORKER_OFFLINE', + `Bridge worker ${args.workerId} is offline`, + ); + } + if ( + (args.requireTenantBinding === true && registration.binding == null) || + (registration.binding != null && + (args.tenantId == null || + args.tenantId.length === 0 || + registration.binding.tenantId !== args.tenantId)) + ) { + throw new BridgeStoreError( + 'WORKER_UNAUTHORIZED', + `Bridge worker ${args.workerId} is not authorized for this tenant`, + ); + } + if ( + args.runtimeSessionId !== undefined && + registration.capabilities.statefulWorkspace !== true + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} does not provide a stateful workspace`, + ); + } + if ( + args.runtimeSessionId !== undefined && + (await this.dispatchCommand( + () => + this.redis.exists( + workspaceQuarantineKey(args.workerId, args.runtimeSessionId ?? ''), + ), + args, + 'Bridge workspace fence read', + )) === 1 + ) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace is quarantined after an incomplete result commit', + ); + } + + const assignmentId = randomBytes(18).toString('base64url'); + const leaseToken = randomBytes(32).toString('base64url'); + const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs); + const 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.dispatchCommand( + () => this.redis.incr(generationKey(args.workerId)), + args, + 'Bridge assignment generation allocation', + ); + assignment = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + assignmentId, + workerId: args.workerId, + incarnationId: registration.incarnationId, + generation, + leaseToken, + leaseTokenHash: tokenHash(leaseToken), + ...(registration.identityId != null + ? { workerIdentityId: registration.identityId } + : {}), + expiresAt: new Date(args.deadlineAtMs).toISOString(), + runtimeSessionId: args.runtimeSessionId, + request: { + body: args.body, + headers: args.headers, + }, + }; + let queued = false; + for (let attempt = 0; attempt < 8 && !queued; attempt += 1) { + this.assertDispatchActive(args.signal, args.deadlineAtMs); + assignment.incarnationId = registration.incarnationId; + queued = await this.dispatchCommand( + () => this.enqueueForActiveIncarnation(assignment!, ttlSeconds), + args, + 'Bridge assignment enqueue', + ); + if (queued) break; + const replacement = await this.dispatchCommand( + () => this.registration(args.workerId), + args, + 'Bridge replacement registration read', + ); + 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, + ); + try { + const result = + args.finalize == null + ? settlement + : await args.finalize(settlement); + await this.commitPendingWorkspace( + assignment, + settlement, + args.deadlineAtMs, + args.signal, + ); + resultCommitted = true; + return result; + } catch (error) { + if (args.runtimeSessionId !== undefined) { + await this.quarantine( + args.workerId, + assignment.incarnationId, + args.runtimeSessionId, + ); + } + throw error; + } + } finally { + if (resultCommitted) { + try { + await this.cleanupWithRetry(args.workerId, assignmentId, assignment); + } catch { + // The lock and assignment have deadline-derived TTLs. Preserve the + // already committed result rather than turning cleanup availability + // into a client-visible failure that could prompt duplicate work. + } + } else { + await this.cleanupDispatch(args.workerId, assignmentId, assignment); + } + } + } + + async lease( + workerId: string, + incarnationId: string, + waitMs: number, + signal?: AbortSignal, + identityId?: string, + ): Promise { + const deadline = Date.now() + waitMs; + let firstPoll = true; + while ( + !signalAborted(signal) && + (firstPoll || Date.now() < deadline) + ) { + firstPoll = false; + let assignmentId: string | null; + try { + assignmentId = await this.leaseCommand( + this.claimOrPopLease(workerId, incarnationId, identityId), + 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())), + signal, + ); + continue; + } + try { + const assignment = await this.leaseCommand( + this.readAssignment(assignmentId), + signal, + 'Bridge lease assignment read', + ); + if ( + assignment == null || + assignment.workerId !== workerId || + assignment.incarnationId !== incarnationId + ) { + 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.leaseCommand( + this.registration(workerId), + signal, + 'Bridge lease registration read', + ); + if (registration?.incarnationId !== incarnationId) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Bridge worker incarnation was replaced', + ); + } + if (assignment.workerIdentityId !== identityId) { + await this.leaseCommand( + this.discardLeaseClaim(workerId, incarnationId, assignmentId), + signal, + 'Bridge unauthorized lease discard', + ); + continue; + } + if (Date.parse(assignment.expiresAt) <= Date.now()) { + 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)) { + await this.returnLease(assignment); + return undefined; + } + const { + leaseTokenHash: _leaseTokenHash, + workerIdentityId: _workerIdentityId, + ...wireAssignment + } = assignment; + return { + ...wireAssignment, + remainingMs: Math.max( + 0, + Date.parse(assignment.expiresAt) - Date.now(), + ), + }; + } catch (error) { + await this.returnLeaseByIdWithRetry( + workerId, + incarnationId, + assignmentId, + ); + if (signalAborted(signal)) return undefined; + throw error; + } + } + return undefined; + } + + async acknowledgeLease( + workerId: string, + incarnationId: string, + assignmentId: string, + generation: number, + leaseToken: string, + signal?: AbortSignal, + ): Promise { + 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 || + 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.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) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment is not the active lease claim', + ); + } + } + + private async claimOrPopLease( + workerId: string, + incarnationId: string, + identityId?: string, + ): Promise { + const result = await this.redis.eval( + [ + "if ARGV[1] ~= '' then", + " if redis.call('GET', KEYS[3]) ~= ARGV[1] then return nil end", + "elseif redis.call('EXISTS', KEYS[3]) == 1 then", + ' return nil', + 'end', + "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'), + 3, + queueKey(workerId, incarnationId), + leaseClaimKey(workerId, incarnationId), + workerStableIdentityKey(workerId), + identityId ?? '', + ); + 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, + assignment.incarnationId, + assignment.assignmentId, + ); + } + + private async returnLeaseById( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { + 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', + ); + } + + 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; + } + + 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, + settlement: CodeBridgeSettlement, + signal?: AbortSignal, + identityId?: string, + ): Promise { + const serializedSettlement = JSON.stringify(settlement); + const existingSettlement = await this.leaseCommand( + this.redis.get(settlementKey(assignmentId)), + signal, + 'Bridge settlement existing read', + ); + 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.leaseCommand( + this.readAssignment(assignmentId), + signal, + 'Bridge settlement assignment read', + ); + if (assignment == null) { + throw new BridgeStoreError( + 'ASSIGNMENT_NOT_FOUND', + 'Bridge assignment was not found', + ); + } + if (assignment.workerId !== workerId) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Bridge assignment belongs to another worker', + ); + } + const registration = await this.leaseCommand( + this.registration(workerId), + signal, + 'Bridge settlement registration read', + ); + if ( + settlement.incarnationId !== assignment.incarnationId || + registration?.incarnationId !== settlement.incarnationId || + settlement.generation !== assignment.generation || + tokenHash(settlement.leaseToken) !== assignment.leaseTokenHash || + assignment.workerIdentityId !== identityId + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment lease is stale', + ); + } + if ( + settlement.status !== 'rejected' && + Date.parse(assignment.expiresAt) <= Date.now() + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment has expired', + ); + } + const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); + const settlementKeys = [ + assignmentKey(assignmentId), + settlementKey(assignmentId), + leaseClaimKey(workerId, assignment.incarnationId), + leaseAckKey(workerId, assignment.incarnationId), + assignmentDeadlineKey(assignmentId), + ]; + if (assignment.runtimeSessionId !== undefined) { + settlementKeys.push( + workspaceQuarantineKey(workerId, assignment.runtimeSessionId), + ); + } + const hasWorkspace = assignment.runtimeSessionId !== undefined; + settlementKeys.push( + `${PREFIX}:stable-identity:${workerId}`, + workerIncarnationKey(workerId), + ); + 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 ARGV[6] == "1" 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', + 'local stableIdentityKey = KEYS[#KEYS - 1]', + 'if ARGV[5] ~= "" then', + ' if redis.call(\'GET\', stableIdentityKey) ~= ARGV[5] then return -4 end', + 'elseif redis.call(\'EXISTS\', stableIdentityKey) == 1 then return -4', + 'end', + 'if redis.call(\'GET\', KEYS[#KEYS]) ~= ARGV[7] then return -4 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 ARGV[6] == "1" and ARGV[4] == "rejected" then redis.call(\'DEL\', KEYS[6]) end', + 'return 1', + ].join('\n'); + const accepted = Number( + await this.leaseCommand( + this.redis.eval( + script, + settlementKeys.length, + ...settlementKeys, + serializedSettlement, + String(ttlSeconds), + assignmentId, + settlement.status, + identityId ?? '', + hasWorkspace ? '1' : '0', + settlement.incarnationId, + ), + signal, + 'Bridge settlement commit', + ), + ); + 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 === -3) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment expired before settlement was committed', + ); + } + if (accepted === -4) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment owner changed before settlement was committed', + ); + } + if (accepted !== 1 && accepted !== 2) { + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment closed before settlement was committed', + ); + } + } + + async cancelled( + workerId: string, + incarnationId: string, + assignmentId: string, + signal?: AbortSignal, + ): Promise { + 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 || + assignment.incarnationId !== incarnationId || + registration?.incarnationId !== incarnationId + ) { + return true; + } + return ( + (await this.leaseCommand( + this.redis.exists(cancellationKey(assignmentId)), + signal, + 'Bridge cancellation marker read', + )) === 1 + ); + } + + async quarantine( + workerId: string, + incarnationId: string, + runtimeSessionId?: string, + ): Promise { + const script = [ + 'redis.call(\'SET\', KEYS[2], \"1\")', + 'if #KEYS == 4 then redis.call(\'SET\', KEYS[4], \"1\") end', + 'local current = redis.call(\'GET\', KEYS[3])', + 'if current == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1], KEYS[3])', + 'end', + 'return 0', + ].join('\n'); + const keys = [ + workerKey(workerId), + quarantineKey(workerId, incarnationId), + workerIncarnationKey(workerId), + ]; + if (runtimeSessionId !== undefined) { + keys.push(workspaceQuarantineKey(workerId, runtimeSessionId)); + } + await boundedCommand( + this.redis.eval( + script, + keys.length, + ...keys, + incarnationId, + ), + this.redisCommandTimeoutMs, + 'Bridge worker quarantine', + ); + } + + async resetWorkspace( + workerId: string, + incarnationId: string, + runtimeSessionId: string, + signal?: AbortSignal, + ): Promise { + const result = Number( + 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) { + 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 { + const raw = await this.redis.get(workerKey(workerId)); + return raw == null ? undefined : (JSON.parse(raw) as RegisteredBridgeWorker); + } + + 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 { + 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 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); + } + 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, + closeKeys.length, + ...closeKeys, + assignment.assignmentId, + ), + this.redisCommandTimeoutMs, + 'Bridge settlement close', + ); + if (finalSettlement != null) { + return JSON.parse(String(finalSettlement)) as CodeBridgeSettlement; + } + throw new BridgeStoreError( + 'ASSIGNMENT_EXPIRED', + 'Bridge assignment exceeded its deadline', + ); + } + + private async cancel( + assignmentId: string, + assignment?: StoredAssignment, + ): Promise { + const ttlSeconds = + assignment == null + ? 30 + : assignmentTtlSeconds(Date.parse(assignment.expiresAt)); + await boundedCommand( + this.redis.set( + cancellationKey(assignmentId), + '1', + 'EX', + ttlSeconds, + ), + this.redisCommandTimeoutMs, + 'Bridge assignment cancellation', + ); + } + + private async enqueueForActiveIncarnation( + assignment: StoredAssignment, + ttlSeconds: number, + ): Promise { + const script = [ + 'if redis.call(\'GET\', KEYS[1]) ~= ARGV[1] then return 0 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])', + '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'); + const keys = [ + workerIncarnationKey(assignment.workerId), + assignmentKey(assignment.assignmentId), + queueKey(assignment.workerId, assignment.incarnationId), + lockIncarnationKey(assignment.workerId), + assignmentDeadlineKey(assignment.assignmentId), + ]; + 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), + String(Date.parse(assignment.expiresAt)), + ); + if (Number(result) === -1) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace already has incomplete stateful work', + ); + } + return Number(result) === 1; + } + + private async acquireLock( + workerId: string, + assignmentId: string, + incarnationId: string, + ttlSeconds: number, + ): Promise { + const script = [ + 'if redis.call(\'EXISTS\', KEYS[1]) == 1 then return 0 end', + 'redis.call(\'SET\', KEYS[1], ARGV[1], \"PX\", ARGV[3])', + 'redis.call(\'SET\', KEYS[2], ARGV[2], \"PX\", ARGV[3])', + 'return 1', + ].join('\n'); + const result = await this.redis.eval( + script, + 2, + lockKey(workerId), + lockIncarnationKey(workerId), + assignmentId, + incarnationId, + String(ttlSeconds * 1000), + ); + return Number(result) === 1; + } + + private async cleanupDispatch( + workerId: string, + assignmentId: string, + assignment: StoredAssignment | undefined, + ): Promise { + await Promise.all([ + this.cancel(assignmentId, assignment), + assignment == null + ? boundedCommand( + this.releaseLock(workerId, assignmentId), + this.redisCommandTimeoutMs, + 'Bridge assignment lock release', + ) + : this.cleanup(assignment), + ]); + } + + private async commitPendingWorkspace( + assignment: StoredAssignment, + settlement: CodeBridgeSettlement, + deadlineAtMs: number, + signal: AbortSignal, + ): Promise { + if ( + assignment.runtimeSessionId === undefined || + settlement.status !== 'fulfilled' + ) { + return; + } + const runtimeSessionId = assignment.runtimeSessionId; + 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 boundedCommand( + this.redis.eval( + script, + 1, + workspaceQuarantineKey( + assignment.workerId, + runtimeSessionId, + ), + assignment.assignmentId, + ), + Math.max( + 1, + Math.min(this.redisCommandTimeoutMs, deadlineAtMs - Date.now()), + ), + 'Bridge workspace commit', + signal, + ), + ); + if (committed !== 1) { + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Bridge workspace commit marker was lost before finalization completed', + ); + } + } + + 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 { + 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, + ), + ]; + const cleanupScript = [ + "local queued = redis.call('LREM', KEYS[2], 0, ARGV[1])", + "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 acknowledged and ARGV[2] == "1" and redis.call(\'GET\', KEYS[5]) == ARGV[1] then', + ' return -1', + 'end', + "return redis.call('DEL', KEYS[1], KEYS[3], KEYS[4])", + ].join('\n'); + const cleanupResult = Number( + await boundedCommand( + this.redis.eval( + cleanupScript, + keys.length, + ...keys, + assignment.assignmentId, + assignment.runtimeSessionId === undefined ? '0' : '1', + ), + this.redisCommandTimeoutMs, + 'Bridge assignment cleanup', + ), + ); + if (cleanupResult !== -1) { + await boundedCommand( + this.releaseLock(assignment.workerId, assignment.assignmentId), + this.redisCommandTimeoutMs, + 'Bridge assignment lock release', + ); + } + } + + private async releaseLock( + workerId: string, + assignmentId: string, + ): Promise { + const script = [ + 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1], KEYS[2])', + 'end', + 'return 0', + ].join('\n'); + await this.redis.eval( + script, + 2, + lockKey(workerId), + lockIncarnationKey(workerId), + assignmentId, + ); + } +} diff --git a/service/src/config.test.ts b/service/src/config.test.ts index 2f88a87a..87658f98 100644 --- a/service/src/config.test.ts +++ b/service/src/config.test.ts @@ -1,6 +1,7 @@ import { describe, expect, test } from 'bun:test'; import { parsePlanLimits, + resolveBridgeAuthMode, resolveRuntimeSessionMode, resolveSandboxBackend, } from './config'; @@ -9,19 +10,23 @@ describe('sandbox execution configuration', () => { test('defaults only unset backend and session mode values', () => { expect(resolveSandboxBackend(undefined)).toBe('http'); expect(resolveRuntimeSessionMode(undefined)).toBe('stateless'); + expect(resolveBridgeAuthMode(undefined)).toBe('static'); }); 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'); + expect(resolveBridgeAuthMode('static')).toBe('static'); + expect(resolveBridgeAuthMode('paired')).toBe('paired'); }); 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'); @@ -30,6 +35,9 @@ describe('sandbox execution configuration', () => { ); expect(() => resolveRuntimeSessionMode('')).toThrow('CODEAPI_RUNTIME_SESSION_MODE'); expect(() => resolveRuntimeSessionMode(' ')).toThrow('CODEAPI_RUNTIME_SESSION_MODE'); + expect(() => resolveBridgeAuthMode('token')).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE must be one of: static, paired', + ); }); }); diff --git a/service/src/config.ts b/service/src/config.ts index ecf35661..95df4eba 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'], ); } @@ -263,8 +263,20 @@ export function resolveRuntimeSessionMode( ); } +export function resolveBridgeAuthMode( + raw: string | undefined, +): 'static' | 'paired' { + return configuredChoice( + raw, + 'CODEAPI_BRIDGE_AUTH_MODE', + 'static', + ['static', 'paired'], + ); +} + const sandboxBackend = resolveSandboxBackend(process.env.CODEAPI_SANDBOX_BACKEND); const runtimeSessionMode = resolveRuntimeSessionMode(process.env.CODEAPI_RUNTIME_SESSION_MODE); +const bridgeAuthMode = resolveBridgeAuthMode(process.env.CODEAPI_BRIDGE_AUTH_MODE); export const env = { PORT: process.env.SERVICE_PORT ?? 3112, @@ -350,8 +362,17 @@ 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, + /** Permit trusted callers to route each execution to a paired worker ID. */ + BRIDGE_DYNAMIC_WORKERS: process.env.CODEAPI_BRIDGE_DYNAMIC_WORKERS === 'true', + /** Outbound worker selected by the remote-bridge backend. */ + BRIDGE_WORKER_ID: process.env.CODEAPI_BRIDGE_WORKER_ID ?? '', + /** Static compatibility auth or short-lived proof-of-possession credentials. */ + BRIDGE_AUTH_MODE: bridgeAuthMode, + /** 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/execution-profile.test.ts b/service/src/execution-profile.test.ts index da20bce2..7ba4640f 100644 --- a/service/src/execution-profile.test.ts +++ b/service/src/execution-profile.test.ts @@ -2,8 +2,11 @@ import { describe, expect, test } from 'bun:test'; import { checkExecutionProfileExpectation, queueNamesForExecutionProfile, + queueNameForExecution, resolveExecutionProfile, resolveExecutionProfileSource, + resolveQueuedSandboxBackend, + validateQueuedSandboxBackend, validateQueuedExecutionProfile, } from './execution-profile'; @@ -55,6 +58,42 @@ describe('execution profile queue isolation', () => { other: 'stateful-other-queue', }); }); + + test('routes a persisted remote bridge replay to the bridge queue on a lambda API', () => { + expect( + queueNameForExecution( + 'python', + 'stateful', + 'explicit', + 'remote-bridge', + ), + ).toBe('remote-bridge-python-queue'); + }); + + test('isolates outbound bridge jobs from Lambda consumers', () => { + expect( + queueNamesForExecutionProfile('stateful', 'explicit', 'remote-bridge'), + ).toEqual({ + python: 'remote-bridge-python-queue', + other: 'remote-bridge-other-queue', + }); + expect( + queueNamesForExecutionProfile('stateful', 'explicit', 'lambda-microvm'), + ).toEqual({ + python: 'stateful-python-queue', + other: 'stateful-other-queue', + }); + }); + + test('labels API-only stateful jobs with their Lambda worker backend', () => { + expect(resolveQueuedSandboxBackend('stateful', 'http')).toBe('lambda-microvm'); + expect(resolveQueuedSandboxBackend('default', 'http', 'explicit')).toBe('http'); + expect(resolveQueuedSandboxBackend('stateful', 'remote-bridge')).toBe('remote-bridge'); + }); + + test('leaves the backend unfenced for inferred stateless legacy queues', () => { + expect(resolveQueuedSandboxBackend('default', 'http', 'inferred')).toBeUndefined(); + }); }); describe('execution profile request assertion', () => { @@ -107,3 +146,29 @@ describe('queued execution profile validation', () => { ); }); }); + +describe('queued sandbox backend validation', () => { + test('accepts matching and legacy jobs', () => { + expect(() => + validateQueuedSandboxBackend('remote-bridge', 'remote-bridge'), + ).not.toThrow(); + expect(() => validateQueuedSandboxBackend(undefined, 'http')).not.toThrow(); + expect(() => + validateQueuedSandboxBackend(undefined, 'remote-bridge', 'legacy-bridge-worker'), + ).not.toThrow(); + }); + + test('rejects invalid and cross-backend jobs', () => { + expect(() => validateQueuedSandboxBackend('invalid', 'http')).toThrow( + 'Queued job has invalid sandbox backend', + ); + expect(() => + validateQueuedSandboxBackend('remote-bridge', 'lambda-microvm'), + ).toThrow( + 'Queued job targets the remote-bridge sandbox backend, but worker serves lambda-microvm', + ); + expect(() => + validateQueuedSandboxBackend(undefined, 'lambda-microvm', 'legacy-bridge-worker'), + ).toThrow('Legacy queued bridge job cannot run on the lambda-microvm sandbox backend'); + }); +}); diff --git a/service/src/execution-profile.ts b/service/src/execution-profile.ts index c4951903..38e9bc8c 100644 --- a/service/src/execution-profile.ts +++ b/service/src/execution-profile.ts @@ -1,4 +1,9 @@ export const EXECUTION_PROFILES = ['default', 'stateful'] as const; +export const SANDBOX_BACKENDS = [ + 'http', + 'lambda-microvm', + 'remote-bridge', +] as const; export type ExecutionProfile = typeof EXECUTION_PROFILES[number]; export type ExecutionProfileSource = 'explicit' | 'inferred'; @@ -11,6 +16,30 @@ export interface ExecutionProfileQueueNames { other: string; } +export type SandboxBackendName = typeof SANDBOX_BACKENDS[number]; + +/** Resolve the backend owned by the queue consumer rather than the API pod. + * Stateful API-only pods intentionally retain the HTTP local default while + * dispatching to Lambda workers. */ +export function resolveQueuedSandboxBackend( + profile: ExecutionProfile, + apiBackend: SandboxBackendName, + source: ExecutionProfileSource = 'explicit', +): SandboxBackendName | undefined { + if (profile === 'stateful' && apiBackend === 'http') { + return 'lambda-microvm'; + } + /* An inferred default profile still uses the pre-fencing legacy queues. + * Its API-only process cannot distinguish the supported HTTP and Lambda + * consumers because Lambda-only configuration belongs to the worker pod. + * Preserve that rollout topology by leaving the backend absent, exactly as + * pre-fencing producers did; explicit profiles regain strict fencing. */ + if (profile === 'default' && source === 'inferred' && apiBackend === 'http') { + return undefined; + } + return apiBackend; +} + export function resolveExecutionProfile( raw: string | undefined, runtimeSessionMode: 'stateless' | 'affinity' | 'strict', @@ -54,10 +83,17 @@ const EXPLICIT_PROFILE_QUEUE_NAMES: Record { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); + validateApiBridgePolicy(); validateExecutionProfilePolicy({ requireBackendMatch: false }); - /* No validateSandboxBackendPolicy() here: an API-only pod authenticates and + validateApiSandboxBackendPolicy(); + /* No full 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. */ @@ -151,6 +162,7 @@ async function gracefulStartup(): Promise { validateWorkerHardenedConfig(); validateExecutionProfilePolicy(); validateSandboxBackendPolicy(); + validateApiBridgePolicy(); await validateLifecycleAuthConfig(); configureProfileMetrics(); @@ -249,12 +261,7 @@ export async function gracefulShutdown(): Promise { } // Close queue connections (both API and Worker need this) - await Promise.all([ - pyQueue.close(), - otherQueue.close(), - pyQueueEvents.close(), - otherQueueEvents.close() - ]); + await closeQueueConnections(); logger.info('Queue connections closed'); // Only disconnect Redis if explicitly requested diff --git a/service/src/local-api.ts b/service/src/local-api.ts index 701270d7..df35cb5e 100644 --- a/service/src/local-api.ts +++ b/service/src/local-api.ts @@ -10,17 +10,22 @@ import express, { json, Router } from 'express'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; +import bridgeRouter from './bridge'; import { requestErrorLogger, requestNotFoundLogger } from './middleware/request-error-logger'; import { executionProfileMiddleware } from './middleware/execution-profile'; import { localAuth } from './auth/local'; -import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from './queue'; +import { pyQueue, otherQueue, connection, closeQueueConnections } from './queue'; import { setStartupComplete } from './lifecycle'; // Workers are imported to ensure they're started with the process 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(); @@ -45,6 +50,7 @@ app.get('/v1/health', async (_, res) => { } }); +v1.use('/bridge', bridgeRouter); v1.use(localAuth); v1.use(serviceRouter); v1.use(programmaticRouter); @@ -56,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, @@ -90,12 +98,7 @@ async function localShutdown(): Promise { localShuttingDown = true; logger.info('Shutting down local server...'); try { - await Promise.all([ - pyQueue.close(), - otherQueue.close(), - pyQueueEvents.close(), - otherQueueEvents.close() - ]); + await closeQueueConnections(); try { await shutdownTelemetry(); } catch (telemetryError) { diff --git a/service/src/queue.ts b/service/src/queue.ts index 91f97c36..54fea308 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -7,7 +7,15 @@ import type * as tls from 'tls'; import type * as t from './types'; import { Jobs } from './enum'; import { env } from './config'; -import { queueNamesForExecutionProfile } from './execution-profile'; +import { + queueNameForExecution, + queueNamesForExecutionProfile, +} from './execution-profile'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from './execution-profile'; import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; @@ -60,18 +68,52 @@ const connection = new IORedis({ const queueNames = queueNamesForExecutionProfile( env.EXECUTION_PROFILE, env.EXECUTION_PROFILE_SOURCE, + env.SANDBOX_BACKEND, ); -const pyQueue = new Queue(queueNames.python, { connection }); -const otherQueue = new Queue(queueNames.other, { connection }); +export interface QueueBinding { + queue: Queue; + events: QueueEvents; + language: 'python' | 'bash'; +} + +const queueResources = new Map< + string, + { queue: Queue; events: QueueEvents } +>(); -const pyQueueEvents = new QueueEvents(queueNames.python, { connection }); -const otherQueueEvents = new QueueEvents(queueNames.other, { connection }); +function getQueueResources( + name: string, +): { queue: Queue; events: QueueEvents } { + const existing = queueResources.get(name); + if (existing != null) return existing; + + const queue = new Queue(name, { connection }); + const events = new QueueEvents(name, { connection }); + setMaxListeners(0, queue, events); + const resources = { queue, events }; + queueResources.set(name, resources); + return resources; +} + +export function getExecutionQueueBinding( + language: 'python' | 'bash', + backend: SandboxBackendName | undefined = env.SANDBOX_BACKEND, + profile: ExecutionProfile = env.EXECUTION_PROFILE, + source: ExecutionProfileSource = env.EXECUTION_PROFILE_SOURCE, +): QueueBinding { + const name = queueNameForExecution( + language, + profile, + source, + backend, + ); + return { ...getQueueResources(name), language }; +} + +const { queue: pyQueue, events: pyQueueEvents } = getQueueResources(queueNames.python); +const { queue: otherQueue, events: otherQueueEvents } = getQueueResources(queueNames.other); const queueMetricStates = ['waiting', 'active', 'delayed'] as const; -const queueMetricSources = [ - { name: queueNames.python, queue: pyQueue }, - { name: queueNames.other, queue: otherQueue }, -] as const; const QUEUE_METRICS_TIMEOUT_MS = 1000; async function withTimeout(promise: Promise, timeoutMs: number, message: string): Promise { @@ -90,7 +132,7 @@ async function withTimeout(promise: Promise, timeoutMs: number, message: s } registerBullmqQueueMetricsCollector(async () => { - await Promise.all(queueMetricSources.map(async ({ name, queue }) => { + await Promise.all([...queueResources.entries()].map(async ([name, { queue }]) => { try { const counts = await withTimeout( queue.getJobCounts(...queueMetricStates), @@ -116,4 +158,13 @@ registerBullmqQueueMetricsCollector(async () => { * BullMQ coordination objects. */ setMaxListeners(0, pyQueue, otherQueue, pyQueueEvents, otherQueueEvents); +export async function closeQueueConnections(): Promise { + await Promise.all( + [...queueResources.values()].flatMap(({ queue, events }) => [ + queue.close(), + events.close(), + ]), + ); +} + export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, queueNames, connection }; 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.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts new file mode 100644 index 00000000..c1b9e2c9 --- /dev/null +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -0,0 +1,153 @@ +import { describe, expect, test } from 'bun:test'; + +import type { SandboxExecuteContext, SandboxTransportRequest } from './types'; +import type { RedisBridgeStore } from '../bridge/store'; + +import { BridgeStoreError } from '../bridge/store'; +import { RemoteBridgeSandboxBackend } from './remote-bridge'; + +function request(): SandboxTransportRequest { + return { + body: { language: 'bash' } as never, + headers: {}, + }; +} + +function context(): SandboxExecuteContext { + return { + executionId: 'execution-1', + language: 'bash', + isSynthetic: false, + signal: new AbortController().signal, + tenantId: 'tenant-1', + bridgeWorkerId: 'user-vm', + runtimeSessionMode: 'strict', + }; +} + +describe('RemoteBridgeSandboxBackend', () => { + test('dispatches a dynamically selected worker with a required tenant binding', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2.0', + files: [], + }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend(store, 'default-vm'); + + await expect(backend.execute(request(), context())).resolves.toMatchObject({ + session_id: 'session-1', + }); + expect(dispatched).toMatchObject({ + workerId: 'user-vm', + tenantId: 'tenant-1', + requireTenantBinding: true, + }); + }); + + test('maps tenant authorization rejection to a bridge backend error', async () => { + const store = { + dispatch: async (): ReturnType => { + throw new BridgeStoreError('WORKER_UNAUTHORIZED', 'private tenant detail'); + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend(store, 'default-vm'); + + await expect(backend.execute(request(), context())).rejects.toMatchObject({ + code: 'BRIDGE_WORKER_UNAUTHORIZED', + }); + }); + + test('keeps an explicitly selected singleton on its unbound compatibility route', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2.0', + files: [], + }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend( + store, + 'deployment-worker', + false, + ); + + await backend.execute(request(), { + ...context(), + bridgeWorkerId: 'deployment-worker', + }); + + expect(dispatched).toMatchObject({ + workerId: 'deployment-worker', + requireTenantBinding: false, + }); + }); + + test('requires a binding for the selected default worker in dynamic mode', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2.0', + files: [], + }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend( + store, + 'deployment-worker', + true, + ); + + await backend.execute(request(), { + ...context(), + bridgeWorkerId: 'deployment-worker', + }); + + expect(dispatched).toMatchObject({ + workerId: 'deployment-worker', + requireTenantBinding: true, + }); + }); +}); diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts new file mode 100644 index 00000000..b1a94ae0 --- /dev/null +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -0,0 +1,95 @@ +import type { + SandboxBackend, + SandboxExecuteContext, + SandboxRawResponse, + SandboxTransportRequest, +} from './types'; +import type { RedisBridgeStore } from '../bridge/store'; + +import { env } from '../config'; +import { bridgeStore } from '../bridge'; +import { BridgeStoreError } from '../bridge/store'; +import { SandboxBackendError } from './types'; + +export class RemoteBridgeSandboxBackend implements SandboxBackend { + readonly name = 'remote-bridge' as const; + + constructor( + private readonly store: Pick = bridgeStore, + private readonly workerId: string = env.BRIDGE_WORKER_ID, + private readonly dynamicWorkers: boolean = env.BRIDGE_DYNAMIC_WORKERS, + ) {} + + async execute( + req: SandboxTransportRequest, + ctx: SandboxExecuteContext, + ): Promise { + const workerId = ctx.bridgeWorkerId ?? this.workerId; + if (workerId.length === 0) { + throw new SandboxBackendError( + 'BRIDGE_WORKER_OFFLINE', + 'No bridge worker is configured', + ); + } + const sessionResultFinalizer = ctx.sessionResultFinalizer; + try { + const settlement = await this.store.dispatch({ + workerId, + tenantId: ctx.tenantId, + requireTenantBinding: + ctx.bridgeWorkerId != null && + (this.dynamicWorkers || ctx.bridgeWorkerId !== this.workerId), + body: req.body, + headers: req.headers, + runtimeSessionId: ctx.runtimeSessionId, + deadlineAtMs: ctx.deadlineAtMs ?? Date.now() + env.JOB_TIMEOUT, + signal: ctx.signal, + finalize: sessionResultFinalizer + ? async (settlement) => { + if (settlement.status === 'rejected') return settlement; + return { + ...settlement, + result: await sessionResultFinalizer(settlement.result), + }; + } + : undefined, + }); + if (settlement.status === 'rejected') { + throw new SandboxBackendError( + 'BRIDGE_EXECUTION_FAILED', + settlement.error, + ); + } + return settlement.result as SandboxRawResponse; + } catch (error) { + if (!(error instanceof BridgeStoreError)) throw error; + if (error.code === 'WORKER_UNAUTHORIZED') { + throw new SandboxBackendError( + 'BRIDGE_WORKER_UNAUTHORIZED', + error.message, + 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..64e5b6e0 100644 --- a/service/src/sandbox-backend/types.ts +++ b/service/src/sandbox-backend/types.ts @@ -37,6 +37,8 @@ export interface SandboxExecuteContext { deadlineAtMs?: number; tenantId?: string; canonicalUserId?: string; + /** Trusted API-selected outbound worker. Presence requires a tenant-bound credential. */ + bridgeWorkerId?: string; /** Absent ⇒ stateless execution (no runtime session affinity). */ runtimeSessionId?: string; runtimeSessionMode: t.RuntimeSessionMode; @@ -56,13 +58,18 @@ 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_UNAUTHORIZED' + | 'BRIDGE_WORKER_BUSY' + | 'BRIDGE_EXECUTION_FAILED' + | 'BRIDGE_DEADLINE_EXCEEDED' | 'MICROVM_LAUNCH_FAILED' | 'MICROVM_LAUNCH_THROTTLED' | 'MICROVM_UNHEALTHY' diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 4aa603e4..78a49497 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -1,7 +1,9 @@ import { afterEach, describe, expect, test } from 'bun:test'; import { env } from './config'; import { + validateApiBridgePolicy, validateApiHardenedConfig, + validateApiSandboxBackendPolicy, validateEgressGatewayHardenedConfig, validateExecutionProfilePolicy, validateSandboxBackendPolicy, @@ -14,6 +16,10 @@ const saved = { executionProfile: env.EXECUTION_PROFILE, executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: env.SANDBOX_BACKEND, + bridgeDynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, + bridgeWorkerId: env.BRIDGE_WORKER_ID, + bridgeAuthMode: env.BRIDGE_AUTH_MODE, + bridgeToken: env.BRIDGE_TOKEN, ptcMode: env.PTC_MODE, runtimeSessionMode: env.RUNTIME_SESSION_MODE, lambdaImageArn: env.LAMBDA_MICROVM_IMAGE_ARN, @@ -52,6 +58,10 @@ function restore(): void { env.EXECUTION_PROFILE = saved.executionProfile; env.EXECUTION_PROFILE_SOURCE = saved.executionProfileSource; env.SANDBOX_BACKEND = saved.sandboxBackend; + env.BRIDGE_DYNAMIC_WORKERS = saved.bridgeDynamicWorkers; + env.BRIDGE_WORKER_ID = saved.bridgeWorkerId; + env.BRIDGE_AUTH_MODE = saved.bridgeAuthMode; + 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 +289,174 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); - test('stateful runtime session modes require the lambda backend', () => { + test('stateful runtime session modes require a stateful backend', () => { env.SANDBOX_BACKEND = 'http'; env.RUNTIME_SESSION_MODE = 'affinity'; - expect(() => validateSandboxBackendPolicy()).toThrow('requires the lambda-microvm backend'); + expect(() => validateSandboxBackendPolicy()).toThrow( + 'requires the lambda-microvm or remote-bridge backend', + ); + env.RUNTIME_SESSION_MODE = 'strict'; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'requires the lambda-microvm or remote-bridge backend', + ); + }); + + test('accepts a configured remote bridge and fails closed on missing enrollment', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.RUNTIME_SESSION_MODE = 'strict'; + 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('allows dynamic-only paired workers without a configured default', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; env.RUNTIME_SESSION_MODE = 'strict'; - expect(() => validateSandboxBackendPolicy()).toThrow('requires the lambda-microvm backend'); + env.PTC_MODE = 'replay'; + env.BRIDGE_DYNAMIC_WORKERS = true; + env.BRIDGE_WORKER_ID = ''; + env.BRIDGE_TOKEN = 'development-bridge-token'; + env.BRIDGE_AUTH_MODE = 'static'; + + expect(() => validateSandboxBackendPolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('requires paired dynamic worker auth in an API-only process', () => { + env.SANDBOX_BACKEND = 'http'; + env.BRIDGE_DYNAMIC_WORKERS = true; + env.BRIDGE_AUTH_MODE = 'static'; + + expect(() => validateApiSandboxBackendPolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateApiSandboxBackendPolicy()).not.toThrow(); + }); + + test('hardened remote bridge requires replay PTC, paired auth, and a strong administrator token', () => { + 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()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('API-only hardened bridge validation rejects static worker auth', () => { + env.SANDBOX_BACKEND = 'http'; + env.HARDENED_SANDBOX_MODE = true; + env.BRIDGE_AUTH_MODE = 'static'; + env.BRIDGE_WORKER_ID = 'engineering-vm'; + env.BRIDGE_TOKEN = 'strong-remote-bridge-token-32-bytes'; + expect(() => validateApiBridgePolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateApiBridgePolicy()).not.toThrow(); + + env.BRIDGE_TOKEN = 'guessable'; + expect(() => validateApiBridgePolicy()).toThrow('at least 32 bytes'); + }); + + 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; + env.BRIDGE_AUTH_MODE = 'paired'; + + expect(() => validateApiBridgePolicy()).toThrow('at least 32 bytes'); + + env.BRIDGE_TOKEN = 'strong-remote-bridge-token-32-bytes'; + 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('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('paired API routes require a configured worker on every backend', () => { + env.SANDBOX_BACKEND = 'http'; + env.BRIDGE_AUTH_MODE = 'paired'; + env.BRIDGE_TOKEN = 'development-bridge-token'; + env.BRIDGE_WORKER_ID = ''; + + expect(() => validateApiBridgePolicy()).toThrow( + 'CODEAPI_BRIDGE_WORKER_ID', + ); + }); + + test('hardened API routes reject padded bridge tokens on HTTP backends', () => { + env.SANDBOX_BACKEND = 'http'; + env.HARDENED_SANDBOX_MODE = true; + env.BRIDGE_AUTH_MODE = 'paired'; + env.BRIDGE_WORKER_ID = 'engineering-vm'; + env.BRIDGE_TOKEN = ' strong-remote-bridge-token-32-bytes '; + + 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'; + 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', () => { diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index a82dabf0..c1740790 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) { @@ -53,6 +54,44 @@ 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.BRIDGE_TOKEN !== env.BRIDGE_TOKEN.trim()) { + throw new SecureStartupConfigError( + 'CODEAPI_BRIDGE_TOKEN must not contain surrounding whitespace', + ); + } + const bridgeEnabled = + env.SANDBOX_BACKEND === 'remote-bridge' || + env.BRIDGE_AUTH_MODE === 'paired'; + if (bridgeEnabled) { + if (!env.BRIDGE_DYNAMIC_WORKERS) { + 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', + ); + } + } + requireValue('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); + } + if (env.SANDBOX_BACKEND === 'remote-bridge') { + requireSafeWholeNumber('JOB_TIMEOUT', env.JOB_TIMEOUT, 1); + 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)', + ); + } + } + if (!env.HARDENED_SANDBOX_MODE) return; + requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); + if (env.BRIDGE_AUTH_MODE !== 'paired') { + throw new SecureStartupConfigError( + 'Hardened API deployments require CODEAPI_BRIDGE_AUTH_MODE=paired because bridge routes are always exposed', + ); + } +} + export function validateWorkerHardenedConfig(): void { if (!env.HARDENED_SANDBOX_MODE) return; rejectValue('CODEAPI_EGRESS_GRANT_SECRET', process.env.CODEAPI_EGRESS_GRANT_SECRET); @@ -93,27 +132,44 @@ 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', ); } } +export function validateApiSandboxBackendPolicy(): void { + if (env.BRIDGE_DYNAMIC_WORKERS && env.BRIDGE_AUTH_MODE !== 'paired') { + throw new SecureStartupConfigError( + 'Dynamic remote bridge workers require CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + } +} + /** * Backend-selection policy. Unlike the hardened-mode validators, this runs * unconditionally: a misconfigured backend must never half-start. */ export function validateSandboxBackendPolicy(): void { + validateApiSandboxBackendPolicy(); if (env.RUNTIME_SESSION_MODE !== 'stateless' && env.SANDBOX_BACKEND === 'http') { throw new SecureStartupConfigError( `CODEAPI_RUNTIME_SESSION_MODE=${env.RUNTIME_SESSION_MODE} requires ` - + 'the lambda-microvm backend; use stateless mode with the http backend', + + 'the lambda-microvm or remote-bridge backend; use stateless mode with the http backend', ); } + if (env.SANDBOX_BACKEND === 'remote-bridge') { + validateApiBridgePolicy(); + return; + } if (env.SANDBOX_BACKEND !== 'lambda-microvm') return; const numericConfigError = lambdaMicrovmNumericConfigError(env); diff --git a/service/src/service-api.ts b/service/src/service-api.ts index ec15b1ac..79db08d5 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'; import { connection } from './queue'; import { env } from './config'; import logger from './logger'; @@ -28,6 +29,7 @@ app.get('/v1/health', async (_, res) => { } }); +v1.use('/bridge', bridgeRouter); v1.use(apiKeyAuth); v1.use(serviceRouter); diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index eade27fb..baa350e6 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -2,11 +2,15 @@ import axios from 'axios'; import { nanoid } from 'nanoid'; import { Router } from 'express'; import type { Response } from 'express'; -import type { Queue, QueueEvents } from 'bullmq'; import type * as t from '../types'; import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle'; import { executionLimiter } from '../middleware/limits'; -import { pyQueue, pyQueueEvents, otherQueue, otherQueueEvents, connection } from '../queue'; +import { + pyQueue, + pyQueueEvents, + connection, + getExecutionQueueBinding, +} from '../queue'; import { createProgrammaticPayload, extractPendingFromStdout } from '../preamble'; import { findBashToolNameCollision } from '../preamble-bash'; import type { LCTool } from '../preamble'; @@ -25,6 +29,8 @@ import { } from '../metrics'; import { Jobs } from '../enum'; import { env, jobCompletionWaitTimeoutMs } from '../config'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; +import { publicExecutionFailure } from '../utils'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, @@ -35,7 +41,15 @@ import { import { findUnregisteredToolCall } from '../tool-scope'; import { summarizeRequestedFiles } from '../execution-log'; import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; -import { buildReplayExecutionState } from './programmatic-state'; +import { + buildReplayExecutionState, + resolveReplayStateSandboxBackend, +} from './programmatic-state'; +import { + BridgeWorkerSelectionError, + CODEAPI_BRIDGE_WORKER_HEADER, + resolveBridgeWorkerSelection, +} from '../bridge/selection'; import logger from '../logger'; import { type ExecutionState, @@ -328,19 +342,6 @@ async function waitForExecutionState( // Replay mode helpers // --------------------------------------------------------------------------- -interface QueueBinding { - queue: Queue; - events: QueueEvents; - language: 'python' | 'bash'; -} - -function pickQueue(language: 'python' | 'bash'): QueueBinding { - if (language === 'bash') { - return { queue: otherQueue, events: otherQueueEvents, language: 'bash' }; - } - return { queue: pyQueue, events: pyQueueEvents, language: 'python' }; -} - function buildReplayPayload( req: t.AuthenticatedRequest, state: ExecutionState, @@ -396,7 +397,21 @@ async function runReplayIteration( }); } - const { queue, events, language } = pickQueue(state.language ?? 'python'); + const replayBackend = + state.sandboxBackend ?? + (state.bridgeWorkerId != null + ? 'remote-bridge' + : resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + )); + const { queue, events, language } = getExecutionQueueBinding( + state.language ?? 'python', + replayBackend, + state.executionProfile ?? env.EXECUTION_PROFILE, + state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, + ); const job = await queue.add(Jobs.execute, { code: state.userCode ?? '', userId, @@ -407,7 +422,9 @@ async function runReplayIteration( executionId: state.execution_id, tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, - executionProfile: env.EXECUTION_PROFILE, + executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, + sandboxBackend: replayBackend, + ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, @@ -439,9 +456,10 @@ async function handleReplayInitial( params: { apiKeyId: string; userId: string; + bridgeWorkerId?: string; }, ): Promise { - const { apiKeyId, userId } = params; + const { apiKeyId, userId, bridgeWorkerId } = params; const { code, tools, @@ -560,6 +578,15 @@ async function handleReplayInitial( isPyPlot, timeout, language, + bridgeWorkerId, + executionProfile: env.EXECUTION_PROFILE, + executionProfileSource: env.EXECUTION_PROFILE_SOURCE, + sandboxBackend: resolveReplayStateSandboxBackend({ + executionProfile: env.EXECUTION_PROFILE, + executionProfileSource: env.EXECUTION_PROFILE_SOURCE, + apiSandboxBackend: env.SANDBOX_BACKEND, + bridgeWorkerId, + }), }); /** Replay mode persists the full request (`userCode` + `tools` + `files`) * inside `ExecutionState` so continuations can re-enqueue without the @@ -832,7 +859,8 @@ async function runAndRespond( logger.error('Replay iteration failed', { execution_id: state.execution_id, err }); await cleanupExecution(state.execution_id, 'replay'); if (!isDisconnected()) { - const message = (err as Error).message; + const publicFailure = publicExecutionFailure(err); + const message = publicFailure?.body.message ?? (err as Error).message; res.status(200).json({ status: 'error', error: message !== '' ? message : 'Sandbox execution failed', @@ -1023,6 +1051,26 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR } = req.body as t.ProgrammaticRequestBody; const rawBody = req.body as Record; const requestedLanguage: unknown = rawBody.language ?? rawBody.lang; + let bridgeWorkerId: string | undefined; + if (continuation_token == null || continuation_token === '') { + try { + const bridgeSelection = resolveBridgeWorkerSelection({ + backend: env.SANDBOX_BACKEND, + configuredWorkerId: env.BRIDGE_WORKER_ID, + dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, + requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), + trustedWorkerId: principal.codeWorkerId, + }); + bridgeWorkerId = bridgeSelection?.explicit === true + ? bridgeSelection.workerId + : undefined; + } catch (error) { + if (error instanceof BridgeWorkerSelectionError) { + return res.status(error.status).json({ error: error.message }); + } + throw error; + } + } if ( requestedLanguage !== undefined && @@ -1079,9 +1127,9 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR }); } if (env.PTC_MODE === 'replay') { - return await handleReplayInitial(req, res, { apiKeyId, userId }); + return await handleReplayInitial(req, res, { apiKeyId, userId, bridgeWorkerId }); } - return await handleBlocking(req, res, { apiKeyId, userId }); + return await handleBlocking(req, res, { apiKeyId, userId, bridgeWorkerId }); } catch (err) { logger.error(`[${INSTANCE_ID}] Programmatic routing error:`, err); if (!res.headersSent) { @@ -1099,9 +1147,9 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR async function handleBlocking( req: t.AuthenticatedRequest, res: Response, - params: { apiKeyId: string; userId: string }, + params: { apiKeyId: string; userId: string; bridgeWorkerId?: string }, ): Promise> { - const { apiKeyId, userId } = params; + const { apiKeyId, userId, bridgeWorkerId } = params; const { code, tools, @@ -1282,6 +1330,7 @@ async function handleBlocking( principalSource: identity.principalSource, authContextHash: identity.authContextHash, apiKeyId, + bridgeWorkerId, startTime: Date.now(), lastActivity: Date.now(), mode: 'blocking', @@ -1378,6 +1427,12 @@ async function handleBlocking( tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + ), + ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index f710668a..fc84d8f8 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -1,7 +1,10 @@ import { describe, expect, test } from 'bun:test'; import type { CodeApiAuthContext, RequestFile } from '../types'; import type { LCTool } from '../preamble'; -import { buildReplayExecutionState } from './programmatic-state'; +import { + buildReplayExecutionState, + resolveReplayStateSandboxBackend, +} from './programmatic-state'; const TOOLS = [ { @@ -22,7 +25,9 @@ const FILES = [ }, ] as RequestFile[]; -function build(overrides: Partial[0]> = {}) { +function build( + overrides: Partial[0]> = {}, +): ReturnType { return buildReplayExecutionState({ executionId: 'exec_123', sessionId: 'session_123', @@ -35,12 +40,35 @@ function build(overrides: Partial[0 isPyPlot: false, timeout: 300000, language: 'python', + executionProfile: 'default', + executionProfileSource: 'inferred', now: 1778250000000, ...overrides, }); } describe('buildReplayExecutionState', () => { + test('persists the resolved queue consumer backend for split stateful deployments', () => { + expect( + resolveReplayStateSandboxBackend({ + executionProfile: 'stateful', + executionProfileSource: 'explicit', + apiSandboxBackend: 'http', + }), + ).toBe('lambda-microvm'); + }); + + test('pins bridge replay state to the remote bridge backend', () => { + expect( + resolveReplayStateSandboxBackend({ + executionProfile: 'stateful', + executionProfileSource: 'explicit', + apiSandboxBackend: 'http', + bridgeWorkerId: 'worker-1', + }), + ).toBe('remote-bridge'); + }); + test('persists canonical LibreChat auth context for replay continuations', () => { const authContext: CodeApiAuthContext = { userId: 'user_canonical', @@ -52,7 +80,13 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', }; - const state = build({ authContext }); + const state = build({ + authContext, + bridgeWorkerId: 'code-user_123', + sandboxBackend: 'remote-bridge', + executionProfile: 'stateful', + executionProfileSource: 'explicit', + }); expect(state).toMatchObject({ execution_id: 'exec_123', @@ -67,6 +101,10 @@ describe('buildReplayExecutionState', () => { principalSource: 'openid_reuse', authContextHash: 'hash_123', apiKeyId: 'key_legacy', + bridgeWorkerId: 'code-user_123', + sandboxBackend: 'remote-bridge', + executionProfile: 'stateful', + executionProfileSource: 'explicit', mode: 'replay', userCode: 'print("hello")', tools: TOOLS, diff --git a/service/src/service/programmatic-state.ts b/service/src/service/programmatic-state.ts index f469606f..25571fed 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -2,6 +2,26 @@ import type * as t from '../types'; import type { LCTool } from '../preamble'; import type { ExecutionState } from './replay-state'; import { buildExecutionIdentity, type ExecutionIdentity } from '../execution-identity'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from '../execution-profile'; + +export function resolveReplayStateSandboxBackend(params: { + executionProfile: ExecutionProfile; + executionProfileSource: ExecutionProfileSource; + apiSandboxBackend: SandboxBackendName; + bridgeWorkerId?: string; +}): SandboxBackendName | undefined { + if (params.bridgeWorkerId != null) return 'remote-bridge'; + return resolveQueuedSandboxBackend( + params.executionProfile, + params.apiSandboxBackend, + params.executionProfileSource, + ); +} export interface BuildReplayExecutionStateParams { executionId: string; @@ -17,6 +37,10 @@ export interface BuildReplayExecutionStateParams { isPyPlot: boolean; timeout: number; language: 'python' | 'bash'; + bridgeWorkerId?: string; + sandboxBackend?: SandboxBackendName; + executionProfile: ExecutionProfile; + executionProfileSource: ExecutionProfileSource; now?: number; } @@ -41,6 +65,10 @@ export function buildReplayExecutionState( principalSource: identity.principalSource, authContextHash: identity.authContextHash, apiKeyId: params.apiKeyId, + bridgeWorkerId: params.bridgeWorkerId, + sandboxBackend: params.sandboxBackend, + executionProfile: params.executionProfile, + executionProfileSource: params.executionProfileSource, startTime: now, lastActivity: now, mode: 'replay', diff --git a/service/src/service/replay-state.ts b/service/src/service/replay-state.ts index 562e06dd..4d2a0501 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -24,6 +24,11 @@ import { nanoid } from 'nanoid'; import type { Redis } from 'ioredis'; import type * as t from '../types'; import type { LCTool } from '../preamble'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from '../execution-profile'; import { connection } from '../queue'; import { env } from '../config'; import { internalServiceHeaders } from '../internal-service-auth'; @@ -109,6 +114,14 @@ export interface ExecutionState { * after one `EXECUTION_STATE_TTL` window post a trusted-source * apiKeyId invariant. */ apiKeyId?: string; + /** Authenticated worker selection retained across every replay iteration. */ + bridgeWorkerId?: string; + /** Original queue/backend target retained across replay continuations. */ + sandboxBackend?: SandboxBackendName; + /** Original producer profile retained so continuations use the same queue. */ + executionProfile?: ExecutionProfile; + /** Original profile source retained because inferred profiles use legacy queues. */ + executionProfileSource?: ExecutionProfileSource; startTime: number; /** * Wall-clock ms of the last interaction that advanced this execution (initial diff --git a/service/src/service/router.ts b/service/src/service/router.ts index f355c2dd..542fe4c0 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -25,7 +25,13 @@ import { Jobs, Languages } from '../enum'; import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; import { createUploadSessionRegistrar } from './upload-session'; import { prepareSandboxJobSecurity } from '../sandbox-egress'; +import { + BridgeWorkerSelectionError, + CODEAPI_BRIDGE_WORKER_HEADER, + resolveBridgeWorkerSelection, +} from '../bridge/selection'; import logger from '../logger'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; const { INSTANCE_ID } = env; const JOB_COMPLETION_WAIT_TIMEOUT_MS = jobCompletionWaitTimeoutMs( @@ -140,6 +146,25 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) return res.status(400).json({ error: `Unsupported language: ${rawLang}` }); } + let bridgeWorkerId: string | undefined; + try { + const bridgeSelection = resolveBridgeWorkerSelection({ + backend: env.SANDBOX_BACKEND, + configuredWorkerId: env.BRIDGE_WORKER_ID, + dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, + requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), + trustedWorkerId: principal.codeWorkerId, + }); + bridgeWorkerId = bridgeSelection?.explicit === true + ? bridgeSelection.workerId + : undefined; + } catch (error) { + if (error instanceof BridgeWorkerSelectionError) { + return res.status(error.status).json({ error: error.message }); + } + throw error; + } + let runtimeSessionId: string | undefined; try { runtimeSessionId = resolveRuntimeSessionIdForExecRequest({ @@ -247,6 +272,12 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + ), + ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), ...(runtimeSessionId != null ? { runtimeSessionId } : {}), runtimeSessionMode, executionManifestClaims: sandboxSecurity.executionManifestClaims, diff --git a/service/src/types/service.ts b/service/src/types/service.ts index a642dda3..555056d5 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -3,7 +3,7 @@ import type { Request } from 'express'; import type { ExecutionManifestClaims } from '../execution-manifest'; import type { ExecutionIdentity } from '../execution-identity'; import type { CodeApiPrincipal } from '../auth/principal'; -import type { ExecutionProfile } from '../execution-profile'; +import type { ExecutionProfile, SandboxBackendName } from '../execution-profile'; import { Jobs } from '@/enum/service'; /** @@ -251,8 +251,12 @@ export type JobData = { executionId?: string; tenantId?: string; canonicalUserId?: string; + /** Trusted dynamic outbound worker selection. */ + bridgeWorkerId?: string; /** Producer deployment identity. Optional only for pre-profile queued jobs. */ executionProfile?: ExecutionProfile; + /** Required sandbox transport. Optional only for jobs queued before fencing. */ + sandboxBackend?: SandboxBackendName; /** * Server-derived runtime session identity. Absence is stateless unless * strict mode requires it; explicit exemptions document intentional gaps. diff --git a/service/src/utils.test.ts b/service/src/utils.test.ts index f1952dfc..86e67a9c 100644 --- a/service/src/utils.test.ts +++ b/service/src/utils.test.ts @@ -48,7 +48,7 @@ describe('isValidResourceId (heterogeneous resource identifiers)', () => { expect(isValidResourceId('682f49b90f07376815c38ef2')).toBe(true); }); - test("accepts 17-char `agent_` slug", () => { + test('accepts 17-char `agent_` slug', () => { expect(isValidResourceId('agent_abc12345678')).toBe(true); }); @@ -145,6 +145,68 @@ describe('sandbox error formatting', () => { }); }); + test('maps bridge authorization and availability failures without leaking worker details', () => { + const unauthorized = publicExecutionFailure( + new Error('BRIDGE_WORKER_UNAUTHORIZED: Worker private-vm belongs to tenant-secret'), + ); + expect(unauthorized).toEqual({ + status: 403, + body: { + error: 'bridge_worker_unauthorized', + message: 'Code environment is not authorized for this tenant', + }, + }); + expect(JSON.stringify(unauthorized)).not.toContain('private-vm'); + expect(JSON.stringify(unauthorized)).not.toContain('tenant-secret'); + + expect( + publicExecutionFailure( + new Error('BRIDGE_WORKER_OFFLINE: Worker private-vm has not checked in'), + ), + ).toEqual({ + status: 503, + body: { + error: 'bridge_worker_offline', + message: 'Code environment is offline', + }, + }); + + const cases = [ + ['BRIDGE_WORKER_BUSY', 409, 'Code environment is busy'], + ['BRIDGE_EXECUTION_FAILED', 502, 'Code environment execution failed'], + [ + 'BRIDGE_DEADLINE_EXCEEDED', + 504, + 'Code environment execution timed out', + ], + ] 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 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: 'Code environment 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 3078d6e8..aae2d05e 100644 --- a/service/src/utils.ts +++ b/service/src/utils.ts @@ -128,15 +128,20 @@ 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_]+):/, ); if (backendMatch) { const code = backendMatch[1]; const statuses: Record = { RUNTIME_SESSION_BUSY: 409, + BRIDGE_WORKER_UNAUTHORIZED: 403, + 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 +152,11 @@ 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_UNAUTHORIZED: 'Code environment is not authorized for this tenant', + BRIDGE_WORKER_OFFLINE: 'Code environment is offline', + BRIDGE_WORKER_BUSY: 'Code environment is busy', + BRIDGE_EXECUTION_FAILED: 'Code environment execution failed', + BRIDGE_DEADLINE_EXCEEDED: 'Code environment execution timed out', MICROVM_LAUNCH_FAILED: 'Sandbox launch failed', MICROVM_LAUNCH_THROTTLED: 'Sandbox capacity is temporarily unavailable', MICROVM_UNHEALTHY: 'Sandbox runtime is unavailable', diff --git a/service/src/workers.ts b/service/src/workers.ts index ed2a46dd..d2048dc8 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -17,7 +17,10 @@ import { isSyntheticPrincipalSource } from './auth/synthetic'; import { withSpan, withTraceContext } from './telemetry'; import { workerDeadlineFailure } from './worker-error'; import logger from './logger'; -import { validateQueuedExecutionProfile } from './execution-profile'; +import { + validateQueuedExecutionProfile, + validateQueuedSandboxBackend, +} from './execution-profile'; const { INSTANCE_ID } = env; const WORKER_ID = `${INSTANCE_ID}-${process.pid}`; @@ -38,7 +41,7 @@ async function processJob(job: t.ExecuteJob): Promise { } async function processJobInner(job: t.ExecuteJob): Promise { - const { code, payload, isPyPlot } = job.data; + const { payload, isPyPlot } = job.data; const isSyntheticJob = job.data.isSynthetic === true || isSyntheticPrincipalSource(job.data.principalSource); const language = payload?.language ?? 'unknown'; const endTimer = jobProcessingDuration.startTimer({ language }); @@ -60,6 +63,11 @@ async function processJobInner(job: t.ExecuteJob): Promise { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); } validateQueuedExecutionProfile(job.data.executionProfile, env.EXECUTION_PROFILE); + validateQueuedSandboxBackend( + job.data.sandboxBackend, + env.SANDBOX_BACKEND, + job.data.bridgeWorkerId, + ); let sandboxPayload = payload; let executionManifestClaims = job.data.executionManifestClaims; let egressGrantToken = job.data.egressGrantToken; @@ -139,6 +147,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { deadlineAtMs, tenantId: job.data.tenantId, canonicalUserId: job.data.canonicalUserId, + bridgeWorkerId: job.data.bridgeWorkerId, runtimeSessionId: runtimeSession.runtimeSessionId, runtimeSessionMode: runtimeSession.runtimeSessionMode, /* Stateful backends run this as a commit barrier after user code but diff --git a/service/tsconfig.json b/service/tsconfig.json index c3e88635..13f30872 100644 --- a/service/tsconfig.json +++ b/service/tsconfig.json @@ -13,7 +13,12 @@ "@/*": ["src/*"] } }, - "include": ["src/**/*.ts", "../shared/telemetry-core.ts"], + "include": [ + "src/**/*.ts", + "../shared/telemetry-core.ts", + "../packages/code/src/protocol.ts", + "../packages/code/src/identity.ts" + ], "exclude": [ "node_modules", "**/*.spec.ts", diff --git a/tests/bridge_pairing_rollout.sh b/tests/bridge_pairing_rollout.sh new file mode 100755 index 00000000..c4246db6 --- /dev/null +++ b/tests/bridge_pairing_rollout.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +set -euo pipefail + +values=helm/codeapi/values.yaml +deployment=helm/codeapi/templates/api-deployment.yaml +rollback=helm/codeapi/scripts/safe-pairing-rollback.sh + +if ! grep -A 12 '^api:$' "$values" | grep -q '^ strategy:$'; then + echo 'api.strategy must be configured for pairing-safe rollouts' >&2 + exit 1 +fi +if ! grep -A 2 '^ strategy:$' "$values" | grep -q '^ type: Recreate$'; then + echo 'api.strategy.type must default to Recreate while pre-fence replicas may exist' >&2 + exit 1 +fi +if ! grep -A 3 '^ strategy:$' "$values" | grep -q '^ rollingUpdate: null$'; then + echo 'api.strategy must clear rollingUpdate when switching existing deployments to Recreate' >&2 + exit 1 +fi +if ! grep -q 'toYaml .Values.api.strategy' "$deployment"; then + echo 'the API Deployment must render api.strategy' >&2 + exit 1 +fi +if ! grep -q 'codeapi.librechat.ai/pairing-fence-version: "1"' "$deployment"; then + echo 'the first pairing-fence chart upgrade must revise the API pod template' >&2 + exit 1 +fi +if ! grep -A 8 '^ image:$' "$values" | grep -q '^ pullPolicy: Always$'; then + echo 'the fenced API rollout must pull the current image even when the default tag is mutable' >&2 + exit 1 +fi +if [[ ! -x "$rollback" ]]; then + echo 'the pairing-safe rollback helper must be executable' >&2 + exit 1 +fi +bash -n "$rollback" +if "$rollback" codeapi 1 default --kube-context other >/dev/null 2>&1; then + echo 'rollback must reject a Helm context that differs from the kubectl drain' >&2 + exit 1 +fi +if "$rollback" codeapi 1 default --kubeconfig=/tmp/other >/dev/null 2>&1; then + echo 'rollback must reject a Helm kubeconfig that differs from the kubectl drain' >&2 + exit 1 +fi +if HELM_KUBECONTEXT=other "$rollback" codeapi 1 default >/dev/null 2>&1; then + echo 'rollback must reject a Helm context inherited from the environment' >&2 + exit 1 +fi +if ! grep -q 'delete horizontalpodautoscaler' "$rollback" || + ! grep -q 'scale "$deployment" --replicas=0' "$rollback" || + ! grep -q -- '--for=delete' "$rollback" || + ! grep -q 'create configmap "$rollback_config_map"' "$rollback" || + ! grep -q 'replica_state=' "$rollback" || + ! grep -q 'discover_api_deployments' "$rollback" || + ! grep -q 'list_api_pods' "$rollback" || + ! grep -q '^ drain_api delete$' "$rollback" || + ! grep -q 'recover_interrupted_rollback' "$rollback" || + ! grep -q 'helm rollback' "$rollback"; then + echo 'rollback must record an epoch, remove autoscaling, verify the drain, and fail closed' >&2 + exit 1 +fi +if ! grep -q 'CODEAPI_BRIDGE_PAIRING_ROLLBACK_EPOCH' "$deployment" || + ! grep -q 'optional: true' "$deployment"; then + echo 'the API Deployment must consume the optional rollback epoch' >&2 + exit 1 +fi