diff --git a/.github/workflows/pkg-pr-new.yaml b/.github/workflows/pkg-pr-new.yaml index 260d428ba4..ec32f60e2c 100644 --- a/.github/workflows/pkg-pr-new.yaml +++ b/.github/workflows/pkg-pr-new.yaml @@ -24,4 +24,4 @@ jobs: find /tmp/inspector-repack -name '*.map' -delete tar czf inspector.tar.gz -C /tmp/inspector-repack . rm -rf /tmp/inspector-repack - - run: pnpm dlx pkg-pr-new publish 'shared/typescript/*' 'engine/sdks/typescript/runner/' 'engine/sdks/typescript/runner-protocol/' 'engine/sdks/typescript/envoy-client/' 'engine/sdks/typescript/envoy-protocol/' 'rivetkit-typescript/packages/*' --packageManager pnpm --template './examples/*' + - run: pnpm dlx pkg-pr-new publish 'shared/typescript/*' 'engine/sdks/typescript/envoy-client/' 'engine/sdks/typescript/envoy-protocol/' 'rivetkit-typescript/packages/*' --packageManager pnpm --template './examples/*' diff --git a/CLAUDE.md b/CLAUDE.md index dc270bb2e7..f79f679850 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -20,7 +20,7 @@ The `rivet.gg` domain is deprecated and should never be used in this codebase. **Never modify an existing published `*.bare` runner protocol version unless explicitly asked to do so.** - Add a new versioned schema instead, then migrate `versioned.rs` and related compatibility code to bridge old versions forward. -- When bumping the protocol version, update `PROTOCOL_MK2_VERSION` in `engine/sdks/rust/runner-protocol/src/lib.rs` and `PROTOCOL_VERSION` in `engine/sdks/typescript/runner/src/mod.ts` together. Both must match the latest schema version. +- When bumping the protocol version, update `PROTOCOL_MK2_VERSION` in `engine/packages/runner-protocol/src/lib.rs` and `PROTOCOL_VERSION` in `rivetkit-typescript/packages/engine-runner/src/mod.ts` together. Both must match the latest schema version. **Keep the KV API in sync between the runner protocol and the KV channel protocol.** diff --git a/Cargo.lock b/Cargo.lock index 88113125e0..22db86ad3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -938,6 +938,15 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "convert_case" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "cookie" version = "0.18.1" @@ -1031,6 +1040,16 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctor" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a2785755761f3ddc1492979ce1e48d2c00d09311c39e4466429188f3dd6501" +dependencies = [ + "quote", + "syn 2.0.104", +] + [[package]] name = "curve25519-dalek" version = "4.1.3" @@ -2849,6 +2868,64 @@ dependencies = [ "vbare", ] +[[package]] +name = "napi" +version = "2.16.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55740c4ae1d8696773c78fdafd5d0e5fe9bc9f1b071c7ba493ba5c413a9184f3" +dependencies = [ + "bitflags", + "ctor", + "napi-derive", + "napi-sys", + "once_cell", + "tokio", +] + +[[package]] +name = "napi-build" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d376940fd5b723c6893cd1ee3f33abbfd86acb1cd1ec079f3ab04a2a3bc4d3b1" + +[[package]] +name = "napi-derive" +version = "2.16.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cbe2585d8ac223f7d34f13701434b9d5f4eb9c332cccce8dee57ea18ab8ab0c" +dependencies = [ + "cfg-if", + "convert_case", + "napi-derive-backend", + "proc-macro2", + "quote", + "syn 2.0.104", +] + +[[package]] +name = "napi-derive-backend" +version = "1.0.75" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1639aaa9eeb76e91c6ae66da8ce3e89e921cd3885e99ec85f4abacae72fc91bf" +dependencies = [ + "convert_case", + "once_cell", + "proc-macro2", + "quote", + "regex", + "semver", + "syn 2.0.104", +] + +[[package]] +name = "napi-sys" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "427802e8ec3a734331fec1035594a210ce1ff4dc5bc1950530920ab717964ea3" +dependencies = [ + "libloading", +] + [[package]] name = "native-tls" version = "0.2.14" @@ -4567,8 +4644,6 @@ dependencies = [ "rivet-cache", "rivet-cache-purge", "rivet-config", - "rivet-engine-runner", - "rivet-envoy-client", "rivet-envoy-protocol", "rivet-guard", "rivet-logs", @@ -4579,6 +4654,7 @@ dependencies = [ "rivet-telemetry", "rivet-term", "rivet-test-deps", + "rivet-test-envoy", "rivet-tracing-reconfigure", "rivet-types", "rivet-util", @@ -4607,23 +4683,13 @@ dependencies = [ ] [[package]] -name = "rivet-engine-runner" +name = "rivet-engine-test-envoy-native" version = "2.2.1" dependencies = [ - "anyhow", - "async-trait", - "chrono", - "futures-util", - "rand 0.8.5", - "rivet-runner-protocol", - "serde", - "serde_bare", - "serde_json", - "tokio", - "tokio-tungstenite", - "tracing", - "urlencoding", - "vbare", + "napi", + "napi-build", + "napi-derive", + "rivet-test-envoy", ] [[package]] @@ -4634,27 +4700,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "rivet-envoy-client" -version = "2.2.1" -dependencies = [ - "anyhow", - "async-trait", - "chrono", - "futures-util", - "rand 0.8.5", - "rivet-envoy-protocol", - "serde", - "serde_bare", - "serde_json", - "tokio", - "tokio-tungstenite", - "tracing", - "urlencoding", - "uuid", - "vbare", -] - [[package]] name = "rivet-envoy-protocol" version = "2.2.1" @@ -4999,6 +5044,34 @@ dependencies = [ "uuid", ] +[[package]] +name = "rivet-test-envoy" +version = "2.2.1" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "axum 0.8.4", + "chrono", + "futures-util", + "hex", + "reqwest", + "rivet-envoy-protocol", + "rivet-runner-protocol", + "rivet-util", + "serde", + "serde_bare", + "serde_json", + "tokio", + "tokio-stream", + "tokio-tungstenite", + "tracing", + "tracing-subscriber", + "urlencoding", + "uuid", + "vbare", +] + [[package]] name = "rivet-tracing-reconfigure" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 96d56d0b6b..616ecde13a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,12 +51,12 @@ members = [ "engine/packages/workflow-worker", "engine/sdks/rust/api-full", "engine/sdks/rust/data", - "engine/sdks/rust/engine-runner", - "engine/sdks/rust/envoy-client", "engine/sdks/rust/envoy-protocol", "engine/sdks/rust/epoxy-protocol", "engine/sdks/rust/kv-channel-protocol", - "engine/sdks/rust/runner-protocol", + "engine/packages/runner-protocol", + "engine/sdks/rust/test-envoy", + "engine/sdks/typescript/test-envoy-native", "engine/sdks/rust/ups-protocol" ] @@ -509,17 +509,14 @@ members = [ [workspace.dependencies.rivet-kv-channel-protocol] path = "engine/sdks/rust/kv-channel-protocol" - [workspace.dependencies.rivet-engine-runner] - path = "engine/sdks/rust/engine-runner" - - [workspace.dependencies.rivet-envoy-client] - path = "engine/sdks/rust/envoy-client" - [workspace.dependencies.rivet-envoy-protocol] path = "engine/sdks/rust/envoy-protocol" [workspace.dependencies.rivet-runner-protocol] - path = "engine/sdks/rust/runner-protocol" + path = "engine/packages/runner-protocol" + + [workspace.dependencies.rivet-test-envoy] + path = "engine/sdks/rust/test-envoy" [workspace.dependencies.rivet-ups-protocol] path = "engine/sdks/rust/ups-protocol" diff --git a/engine/CLAUDE.md b/engine/CLAUDE.md index 19175eb9a4..b3f912c407 100644 --- a/engine/CLAUDE.md +++ b/engine/CLAUDE.md @@ -23,9 +23,9 @@ When changing a versioned VBARE schema, follow the existing migration pattern. 3. Verify the affected Rust crate still builds. 4. For the runner protocol specifically: - Bump both protocol constants together: - - `engine/sdks/rust/runner-protocol/src/lib.rs` `PROTOCOL_MK2_VERSION` - - `engine/sdks/typescript/runner/src/mod.ts` `PROTOCOL_VERSION` - - Update the Rust latest re-export in `engine/sdks/rust/runner-protocol/src/lib.rs` to the new generated module. + - `engine/packages/runner-protocol/src/lib.rs` `PROTOCOL_MK2_VERSION` + - `rivetkit-typescript/packages/engine-runner/src/mod.ts` `PROTOCOL_VERSION` + - Update the Rust latest re-export in `engine/packages/runner-protocol/src/lib.rs` to the new generated module. ## Epoxy durable keys diff --git a/engine/docker/dev-host/docker-compose.yml b/engine/docker/dev-host/docker-compose.yml index f51586ef6b..096604c5bd 100644 --- a/engine/docker/dev-host/docker-compose.yml +++ b/engine/docker/dev-host/docker-compose.yml @@ -180,12 +180,15 @@ services: runner: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://127.0.0.1:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine: diff --git a/engine/docker/dev-host/rivet-engine/config.jsonc b/engine/docker/dev-host/rivet-engine/config.jsonc index 814d17ecbb..d5a3095269 100644 --- a/engine/docker/dev-host/rivet-engine/config.jsonc +++ b/engine/docker/dev-host/rivet-engine/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "default", + "datacenters": { + "default": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://127.0.0.1:6421", @@ -24,7 +23,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@127.0.0.1:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc index a61a8d52f1..ab9cfa618a 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/0/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc index a61a8d52f1..ab9cfa618a 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/1/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc index a61a8d52f1..ab9cfa618a 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-a/rivet-engine/2/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc index 7898758e89..2a34bc83ff 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/0/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 2, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc index 7898758e89..2a34bc83ff 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/1/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 2, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc index 7898758e89..2a34bc83ff 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-b/rivet-engine/2/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 2, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc index 4d40c5693d..970454f7f6 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/0/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 3, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc index 4d40c5693d..970454f7f6 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/1/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 3, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc b/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc index 4d40c5693d..970454f7f6 100644 --- a/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc +++ b/engine/docker/dev-multidc-multinode/datacenters/dc-c/rivet-engine/2/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 3, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a-0:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b-0:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c-0:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" diff --git a/engine/docker/dev-multidc-multinode/docker-compose.yml b/engine/docker/dev-multidc-multinode/docker-compose.yml index 9677c5c437..d1061b6ac5 100644 --- a/engine/docker/dev-multidc-multinode/docker-compose.yml +++ b/engine/docker/dev-multidc-multinode/docker-compose.yml @@ -300,12 +300,15 @@ services: runner-dc-a-0: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s ports: - '5050:5050' @@ -317,12 +320,15 @@ services: runner-dc-a-1: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-a-0: @@ -332,12 +338,15 @@ services: runner-dc-a-2: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-a-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-a-0: @@ -555,12 +564,15 @@ services: runner-dc-b-0: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: @@ -570,12 +582,15 @@ services: runner-dc-b-1: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: @@ -585,12 +600,15 @@ services: runner-dc-b-2: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-b-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-b-0: @@ -808,12 +826,15 @@ services: runner-dc-c-0: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: @@ -823,12 +844,15 @@ services: runner-dc-c-1: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: @@ -838,12 +862,15 @@ services: runner-dc-c-2: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-c-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-c-0: diff --git a/engine/docker/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc b/engine/docker/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc index 4c9e465d85..f989acaae9 100644 --- a/engine/docker/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc +++ b/engine/docker/dev-multidc/datacenters/dc-a/rivet-engine/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-a:5432/rivet_engine" diff --git a/engine/docker/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc b/engine/docker/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc index f35557210f..fa082ce527 100644 --- a/engine/docker/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc +++ b/engine/docker/dev-multidc/datacenters/dc-b/rivet-engine/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 2, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-b:5432/rivet_engine" diff --git a/engine/docker/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc b/engine/docker/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc index c12f4bdc1d..f3c3c6ae38 100644 --- a/engine/docker/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc +++ b/engine/docker/dev-multidc/datacenters/dc-c/rivet-engine/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 3, - "datacenters": [ - { - "name": "dc-a", + "datacenters": { + "dc-a": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-dc-a:6421", @@ -24,8 +23,7 @@ "localhost" ] }, - { - "name": "dc-b", + "dc-b": { "datacenter_label": 2, "is_leader": false, "peer_url": "http://rivet-engine-dc-b:6421", @@ -36,8 +34,7 @@ "localhost" ] }, - { - "name": "dc-c", + "dc-c": { "datacenter_label": 3, "is_leader": false, "peer_url": "http://rivet-engine-dc-c:6421", @@ -48,7 +45,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres-dc-c:5432/rivet_engine" diff --git a/engine/docker/dev-multidc/docker-compose.yml b/engine/docker/dev-multidc/docker-compose.yml index ce49486879..5d4d0b53da 100644 --- a/engine/docker/dev-multidc/docker-compose.yml +++ b/engine/docker/dev-multidc/docker-compose.yml @@ -216,12 +216,15 @@ services: runner-dc-a: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-a:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s ports: - '5050:5050' @@ -357,12 +360,15 @@ services: runner-dc-b: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-b:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-b: @@ -496,12 +502,15 @@ services: runner-dc-c: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-dc-c:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-dc-c: diff --git a/engine/docker/dev-multinode/docker-compose.yml b/engine/docker/dev-multinode/docker-compose.yml index 3538ead70a..f241838c81 100644 --- a/engine/docker/dev-multinode/docker-compose.yml +++ b/engine/docker/dev-multinode/docker-compose.yml @@ -283,12 +283,15 @@ services: runner-0: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s ports: - '5050:5050' @@ -300,12 +303,15 @@ services: runner-1: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-0: @@ -315,12 +321,15 @@ services: runner-2: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine-0:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s depends_on: rivet-engine-0: diff --git a/engine/docker/dev-multinode/rivet-engine/0/config.jsonc b/engine/docker/dev-multinode/rivet-engine/0/config.jsonc index 31b5ce77fe..655fc83ed7 100644 --- a/engine/docker/dev-multinode/rivet-engine/0/config.jsonc +++ b/engine/docker/dev-multinode/rivet-engine/0/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "default", + "datacenters": { + "default": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-0:6421", @@ -24,7 +23,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" diff --git a/engine/docker/dev-multinode/rivet-engine/1/config.jsonc b/engine/docker/dev-multinode/rivet-engine/1/config.jsonc index 31b5ce77fe..655fc83ed7 100644 --- a/engine/docker/dev-multinode/rivet-engine/1/config.jsonc +++ b/engine/docker/dev-multinode/rivet-engine/1/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "default", + "datacenters": { + "default": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-0:6421", @@ -24,7 +23,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" diff --git a/engine/docker/dev-multinode/rivet-engine/2/config.jsonc b/engine/docker/dev-multinode/rivet-engine/2/config.jsonc index 31b5ce77fe..655fc83ed7 100644 --- a/engine/docker/dev-multinode/rivet-engine/2/config.jsonc +++ b/engine/docker/dev-multinode/rivet-engine/2/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "default", + "datacenters": { + "default": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine-0:6421", @@ -24,7 +23,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" diff --git a/engine/docker/dev/docker-compose.yml b/engine/docker/dev/docker-compose.yml index 6e2b881adc..eb7923449e 100644 --- a/engine/docker/dev/docker-compose.yml +++ b/engine/docker/dev/docker-compose.yml @@ -207,12 +207,15 @@ services: runner: build: context: ../../.. - dockerfile: engine/sdks/typescript/test-runner/Dockerfile + dockerfile: engine/sdks/rust/test-envoy/Dockerfile platform: linux/amd64 restart: unless-stopped environment: - RIVET_ENDPOINT=http://rivet-engine:6420 - - RIVET_RUNNER_TOTAL_SLOTS=1000000 + - INTERNAL_SERVER_PORT=5050 + - RIVET_POOL_NAME=test-envoy + - AUTOSTART_ENVOY=1 + - AUTOCONFIGURE_SERVERLESS=0 stop_grace_period: 4s ports: - '5050:5050' diff --git a/engine/docker/dev/rivet-engine/config.jsonc b/engine/docker/dev/rivet-engine/config.jsonc index 3158c9bf72..74135c72fa 100644 --- a/engine/docker/dev/rivet-engine/config.jsonc +++ b/engine/docker/dev/rivet-engine/config.jsonc @@ -11,9 +11,8 @@ }, "topology": { "datacenter_label": 1, - "datacenters": [ - { - "name": "default", + "datacenters": { + "default": { "datacenter_label": 1, "is_leader": true, "peer_url": "http://rivet-engine:6421", @@ -24,7 +23,7 @@ "localhost" ] } - ] + } }, "postgres": { "url": "postgresql://postgres:postgres@postgres:5432/rivet_engine" diff --git a/engine/docker/template/src/docker-compose.ts b/engine/docker/template/src/docker-compose.ts index 62ad7279e3..dac2b97511 100644 --- a/engine/docker/template/src/docker-compose.ts +++ b/engine/docker/template/src/docker-compose.ts @@ -336,14 +336,15 @@ export function generateDockerCompose(context: TemplateContext) { services[serviceName] = { build: { context: "../../..", - dockerfile: "engine/sdks/typescript/test-runner/Dockerfile", + dockerfile: "engine/sdks/rust/test-envoy/Dockerfile", }, platform: "linux/amd64", restart: "unless-stopped", environment: [ `RIVET_ENDPOINT=http://${context.getServiceHost("rivet-engine", datacenter.name, 0)}:6420`, - `RIVET_RUNNER_TOTAL_SLOTS=1000000`, - `AUTOSTART_RUNNER=1`, + `INTERNAL_SERVER_PORT=5050`, + `RIVET_POOL_NAME=test-envoy`, + `AUTOSTART_ENVOY=1`, `AUTOCONFIGURE_SERVERLESS=0` ], stop_grace_period: "4s", diff --git a/engine/docker/template/src/services/edge/runner.ts b/engine/docker/template/src/services/edge/runner.ts index 31ec3b1521..c04f936dd3 100644 --- a/engine/docker/template/src/services/edge/runner.ts +++ b/engine/docker/template/src/services/edge/runner.ts @@ -1,7 +1,6 @@ import type { TemplateContext } from "../../context"; export function generateRunner(context: TemplateContext) { - // The runner files are shared and located at sdks/typescript/test-runner/ - // We just need to reference them in the docker-compose - // No specific files need to be generated here for now + // The test runner service now uses the Rust test-envoy binary. + // The docker-compose template points at the Rust Dockerfile directly. } diff --git a/engine/packages/engine/Cargo.toml b/engine/packages/engine/Cargo.toml index fb365c2c53..c678ffa68b 100644 --- a/engine/packages/engine/Cargo.toml +++ b/engine/packages/engine/Cargo.toml @@ -69,10 +69,9 @@ portpicker.workspace = true rand.workspace = true rivet-api-public.workspace = true rivet-api-types.workspace = true -rivet-engine-runner.workspace = true -rivet-envoy-client.workspace = true rivet-envoy-protocol.workspace = true rivet-runner-protocol.workspace = true +rivet-test-envoy.workspace = true rivet-test-deps.workspace = true rivet-util.workspace = true rstest.workspace = true diff --git a/engine/packages/engine/tests/common/test_envoy.rs b/engine/packages/engine/tests/common/test_envoy.rs index ecb15c9381..1de3fc4f78 100644 --- a/engine/packages/engine/tests/common/test_envoy.rs +++ b/engine/packages/engine/tests/common/test_envoy.rs @@ -1,6 +1,6 @@ -//! Test envoy client wrapper for engine tests. +//! Test envoy wrapper for engine tests. //! -//! This module provides a `TestEnvoyBuilder` that wraps the standalone `rivet-envoy-client` +//! This module provides a `TestEnvoyBuilder` that wraps the standalone `rivet-test-envoy` //! package, adding test-specific functionality like building from a `TestDatacenter`. use anyhow::Result; @@ -8,7 +8,7 @@ use std::collections::HashMap; use std::sync::Arc; // Re-export everything from the standalone package -pub use rivet_envoy_client::{ +pub use rivet_test_envoy::{ ActorConfig, ActorEvent, ActorLifecycleEvent, ActorStartResult, ActorStopResult, CountingCrashActor, CrashNTimesThenSucceedActor, CrashOnStartActor, CustomActor, CustomActorBuilder, DelayedStartActor, EchoActor, Envoy, EnvoyBuilder, EnvoyConfig, KvRequest, diff --git a/engine/packages/engine/tests/common/test_runner.rs b/engine/packages/engine/tests/common/test_runner.rs index 8d3cba0faf..232e8edd0d 100644 --- a/engine/packages/engine/tests/common/test_runner.rs +++ b/engine/packages/engine/tests/common/test_runner.rs @@ -1,28 +1,204 @@ //! Test runner wrapper for engine tests. //! -//! This module provides a `TestRunnerBuilder` that wraps the standalone `rivet-engine-runner` -//! package, adding test-specific functionality like building from a `TestDatacenter`. +//! This module now adapts the Rust `rivet-test-envoy` harness to the legacy +//! runner-oriented test surface that the engine tests still import. use anyhow::Result; use std::collections::HashMap; use std::sync::Arc; -// Re-export everything from the standalone package -pub use rivet_engine_runner::{ +pub use rivet_envoy_protocol::PROTOCOL_VERSION; +pub use rivet_runner_protocol as protocol_types; +pub use rivet_test_envoy::{ ActorConfig, ActorEvent, ActorLifecycleEvent, ActorStartResult, ActorStopResult, CountingCrashActor, CrashNTimesThenSucceedActor, CrashOnStartActor, CustomActor, - CustomActorBuilder, DelayedStartActor, EchoActor, KvRequest, NotifyOnStartActor, - PROTOCOL_VERSION, Runner, RunnerBuilder, RunnerBuilderLegacy, RunnerConfig, - SleepImmediatelyActor, StopImmediatelyActor, TestActor, TimeoutActor, VerifyInputActor, - protocol_types, + CustomActorBuilder, DelayedStartActor, EchoActor, Envoy, EnvoyBuilder as TestEnvoyBuilderImpl, + EnvoyConfig, KvRequest, NotifyOnStartActor, SleepImmediatelyActor, StopImmediatelyActor, + TestActor, TimeoutActor, VerifyInputActor, }; -// Type alias for backwards compatibility +type ActorFactory = Arc Box + Send + Sync>; + pub type TestRunner = Runner; +pub type RunnerBuilderLegacy = RunnerBuilder; -type ActorFactory = Arc Box + Send + Sync>; +#[derive(Clone)] +pub struct RunnerConfig { + endpoint: String, + token: String, + namespace: String, + runner_name: String, + runner_key: String, + version: u32, + total_slots: u32, +} + +impl RunnerConfig { + pub fn builder() -> RunnerConfigBuilder { + RunnerConfigBuilder::default() + } +} + +#[derive(Default)] +pub struct RunnerConfigBuilder { + endpoint: Option, + token: Option, + namespace: Option, + runner_name: Option, + runner_key: Option, + version: Option, + total_slots: Option, +} + +impl RunnerConfigBuilder { + pub fn endpoint(mut self, endpoint: impl Into) -> Self { + self.endpoint = Some(endpoint.into()); + self + } -/// Test-specific runner builder that integrates with TestDatacenter + pub fn token(mut self, token: impl Into) -> Self { + self.token = Some(token.into()); + self + } + + pub fn namespace(mut self, namespace: impl Into) -> Self { + self.namespace = Some(namespace.into()); + self + } + + pub fn runner_name(mut self, runner_name: impl Into) -> Self { + self.runner_name = Some(runner_name.into()); + self + } + + pub fn runner_key(mut self, runner_key: impl Into) -> Self { + self.runner_key = Some(runner_key.into()); + self + } + + pub fn version(mut self, version: u32) -> Self { + self.version = Some(version); + self + } + + pub fn total_slots(mut self, total_slots: u32) -> Self { + self.total_slots = Some(total_slots); + self + } + + pub fn build(self) -> Result { + Ok(RunnerConfig { + endpoint: self + .endpoint + .ok_or_else(|| anyhow::anyhow!("endpoint is required"))?, + token: self.token.unwrap_or_else(|| "dev".to_string()), + namespace: self + .namespace + .ok_or_else(|| anyhow::anyhow!("namespace is required"))?, + runner_name: self + .runner_name + .unwrap_or_else(|| "test-runner".to_string()), + runner_key: self + .runner_key + .unwrap_or_else(|| format!("key-{:012x}", rand::random::())), + version: self.version.unwrap_or(1), + total_slots: self.total_slots.unwrap_or(100), + }) + } +} + +pub struct RunnerBuilder { + config: RunnerConfig, + actor_factories: HashMap, +} + +impl RunnerBuilder { + pub fn new(config: RunnerConfig) -> Self { + Self { + config, + actor_factories: HashMap::new(), + } + } + + pub fn with_actor_behavior(mut self, actor_name: &str, factory: F) -> Self + where + F: Fn(ActorConfig) -> Box + Send + Sync + 'static, + { + self.actor_factories + .insert(actor_name.to_string(), Arc::new(factory)); + self + } + + pub fn build(self) -> Result { + let envoy_config = EnvoyConfig::builder() + .endpoint(&self.config.endpoint) + .token(&self.config.token) + .namespace(&self.config.namespace) + .pool_name(&self.config.runner_name) + .version(self.config.version) + .metadata(serde_json::json!({ + "runner_key": self.config.runner_key, + "total_slots": self.config.total_slots, + })) + .build()?; + + let mut builder = TestEnvoyBuilderImpl::new(envoy_config); + for (name, factory) in self.actor_factories { + builder = builder.with_actor_behavior(&name, move |config| factory(config)); + } + + Ok(Runner { + runner_id: format!("runner-{}", uuid::Uuid::new_v4()), + runner_name: self.config.runner_name, + envoy: builder.build()?, + }) + } +} + +pub struct Runner { + pub runner_id: String, + runner_name: String, + envoy: Envoy, +} + +impl Runner { + pub async fn start(&self) -> Result<()> { + self.envoy.start().await + } + + pub async fn wait_ready(&self) -> String { + self.envoy.wait_ready().await; + self.runner_id.clone() + } + + pub async fn has_actor(&self, actor_id: &str) -> bool { + self.envoy.has_actor(actor_id).await + } + + pub async fn get_actor_ids(&self) -> Vec { + self.envoy.get_actor_ids().await + } + + pub fn name(&self) -> &str { + &self.runner_name + } + + pub fn subscribe_lifecycle_events( + &self, + ) -> tokio::sync::broadcast::Receiver { + self.envoy.subscribe_lifecycle_events() + } + + pub async fn shutdown(&self) { + self.envoy.shutdown().await; + } + + pub async fn crash(&self) { + self.envoy.crash().await; + } +} + +/// Test-specific runner builder that integrates with TestDatacenter. pub struct TestRunnerBuilder { namespace: String, runner_name: String, @@ -64,7 +240,6 @@ impl TestRunnerBuilder { self } - /// Register an actor factory for a specific actor name pub fn with_actor_behavior(mut self, actor_name: &str, factory: F) -> Self where F: Fn(ActorConfig) -> Box + Send + Sync + 'static, @@ -74,12 +249,10 @@ impl TestRunnerBuilder { self } - /// Build the runner using the TestDatacenter's guard port pub async fn build(self, dc: &super::TestDatacenter) -> Result { let endpoint = format!("http://127.0.0.1:{}", dc.guard_port()); let token = "dev".to_string(); - // Build the config using the new API let config = RunnerConfig::builder() .endpoint(&endpoint) .token(&token) @@ -90,10 +263,7 @@ impl TestRunnerBuilder { .total_slots(self.total_slots) .build()?; - // Build the runner let mut builder = RunnerBuilder::new(config); - - // Register all actor factories for (name, factory) in self.actor_factories { builder = builder.with_actor_behavior(&name, move |config| factory(config)); } diff --git a/engine/sdks/rust/runner-protocol/Cargo.toml b/engine/packages/runner-protocol/Cargo.toml similarity index 100% rename from engine/sdks/rust/runner-protocol/Cargo.toml rename to engine/packages/runner-protocol/Cargo.toml diff --git a/engine/sdks/rust/runner-protocol/build.rs b/engine/packages/runner-protocol/build.rs similarity index 85% rename from engine/sdks/rust/runner-protocol/build.rs rename to engine/packages/runner-protocol/build.rs index bf0fd1f848..f39d661e13 100644 --- a/engine/sdks/rust/runner-protocol/build.rs +++ b/engine/packages/runner-protocol/build.rs @@ -6,13 +6,14 @@ use std::{ fn main() -> Result<(), Box> { let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")?; - let workspace_root = Path::new(&manifest_dir) + let repo_root = Path::new(&manifest_dir) .parent() .and_then(|p| p.parent()) .and_then(|p| p.parent()) - .ok_or("Failed to find workspace root")?; + .ok_or("Failed to find repository root")?; + let engine_root = repo_root.join("engine"); - let schema_dir = workspace_root + let schema_dir = engine_root .join("sdks") .join("schemas") .join("runner-protocol"); @@ -22,12 +23,9 @@ fn main() -> Result<(), Box> { vbare_compiler::process_schemas_with_config(&schema_dir, &cfg)?; // TypeScript SDK generation - let cli_js_path = workspace_root - .parent() - .unwrap() - .join("node_modules/@bare-ts/tools/dist/bin/cli.js"); + let cli_js_path = repo_root.join("node_modules/@bare-ts/tools/dist/bin/cli.js"); if cli_js_path.exists() { - typescript::generate_sdk(&schema_dir); + typescript::generate_sdk(repo_root, &schema_dir); } else { println!( "cargo:warning=TypeScript SDK generation skipped: cli.js not found at {}. Run `pnpm install` to install.", @@ -41,18 +39,18 @@ fn main() -> Result<(), Box> { mod typescript { use super::*; - pub fn generate_sdk(schema_dir: &Path) { + pub fn generate_sdk(repo_root: &Path, schema_dir: &Path) { let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); - let workspace_root = Path::new(&manifest_dir) + let _repo_root_from_manifest = Path::new(&manifest_dir) .parent() .and_then(|p| p.parent()) .and_then(|p| p.parent()) - .expect("Failed to find workspace root"); + .expect("Failed to find repository root"); - let sdk_dir = workspace_root - .join("sdks") - .join("typescript") - .join("runner-protocol"); + let sdk_dir = repo_root + .join("rivetkit-typescript") + .join("packages") + .join("engine-runner-protocol"); let src_dir = sdk_dir.join("src"); let highest_version_path = super::find_highest_version(schema_dir); @@ -65,10 +63,7 @@ mod typescript { let output_path = src_dir.join("index.ts"); let output = Command::new( - workspace_root - .parent() - .unwrap() - .join("node_modules/@bare-ts/tools/dist/bin/cli.js"), + repo_root.join("node_modules/@bare-ts/tools/dist/bin/cli.js"), ) .arg("compile") .arg("--generator") diff --git a/engine/sdks/rust/runner-protocol/src/compat.rs b/engine/packages/runner-protocol/src/compat.rs similarity index 100% rename from engine/sdks/rust/runner-protocol/src/compat.rs rename to engine/packages/runner-protocol/src/compat.rs diff --git a/engine/sdks/rust/runner-protocol/src/generated.rs b/engine/packages/runner-protocol/src/generated.rs similarity index 100% rename from engine/sdks/rust/runner-protocol/src/generated.rs rename to engine/packages/runner-protocol/src/generated.rs diff --git a/engine/sdks/rust/runner-protocol/src/lib.rs b/engine/packages/runner-protocol/src/lib.rs similarity index 100% rename from engine/sdks/rust/runner-protocol/src/lib.rs rename to engine/packages/runner-protocol/src/lib.rs diff --git a/engine/sdks/rust/runner-protocol/src/util.rs b/engine/packages/runner-protocol/src/util.rs similarity index 100% rename from engine/sdks/rust/runner-protocol/src/util.rs rename to engine/packages/runner-protocol/src/util.rs diff --git a/engine/sdks/rust/runner-protocol/src/uuid_compat.rs b/engine/packages/runner-protocol/src/uuid_compat.rs similarity index 100% rename from engine/sdks/rust/runner-protocol/src/uuid_compat.rs rename to engine/packages/runner-protocol/src/uuid_compat.rs diff --git a/engine/sdks/rust/runner-protocol/src/versioned.rs b/engine/packages/runner-protocol/src/versioned.rs similarity index 100% rename from engine/sdks/rust/runner-protocol/src/versioned.rs rename to engine/packages/runner-protocol/src/versioned.rs diff --git a/engine/packages/test-deps/src/datacenter.rs b/engine/packages/test-deps/src/datacenter.rs index ede60e9e22..33843062df 100644 --- a/engine/packages/test-deps/src/datacenter.rs +++ b/engine/packages/test-deps/src/datacenter.rs @@ -90,6 +90,10 @@ pub async fn setup_single_datacenter( https: None, ..Default::default() }); + root.metrics = rivet_config::config::metrics::Metrics { + host: None, + port: Some(0), + }; // Use short timeouts for tests root.pegboard = Some(rivet_config::config::pegboard::Pegboard { diff --git a/engine/sdks/rust/engine-runner/Cargo.toml b/engine/sdks/rust/engine-runner/Cargo.toml deleted file mode 100644 index d623eca8a5..0000000000 --- a/engine/sdks/rust/engine-runner/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "rivet-engine-runner" -version.workspace = true -authors.workspace = true -license.workspace = true -edition.workspace = true -description = "Rust-based engine runner for Rivet, enabling programmatic actor lifecycle control" - -[dependencies] -anyhow.workspace = true -async-trait.workspace = true -chrono.workspace = true -futures-util.workspace = true -rand.workspace = true -rivet-runner-protocol.workspace = true -serde_bare.workspace = true -serde_json.workspace = true -serde.workspace = true -tokio.workspace = true -tokio-tungstenite.workspace = true -tracing.workspace = true -urlencoding.workspace = true -vbare.workspace = true diff --git a/engine/sdks/rust/engine-runner/src/actor.rs b/engine/sdks/rust/engine-runner/src/actor.rs deleted file mode 100644 index 029ecec622..0000000000 --- a/engine/sdks/rust/engine-runner/src/actor.rs +++ /dev/null @@ -1,304 +0,0 @@ -use anyhow::Result; -use async_trait::async_trait; -use rivet_runner_protocol::mk2 as rp; -use std::time::Duration; -use tokio::sync::{mpsc, oneshot}; - -use crate::protocol; - -/// Configuration passed to actor when it starts -#[derive(Clone)] -pub struct ActorConfig { - pub actor_id: String, - pub generation: u32, - pub name: String, - pub key: Option, - pub create_ts: i64, - pub input: Option>, - - /// Channel to send events to the runner - pub event_tx: mpsc::UnboundedSender, - - /// Channel to send KV requests to the runner - pub kv_request_tx: mpsc::UnboundedSender, -} - -impl ActorConfig { - pub fn new( - config: &rp::ActorConfig, - actor_id: String, - generation: u32, - event_tx: mpsc::UnboundedSender, - kv_request_tx: mpsc::UnboundedSender, - ) -> Self { - ActorConfig { - actor_id, - generation, - name: config.name.clone(), - key: config.key.clone(), - create_ts: config.create_ts, - input: config.input.as_ref().map(|i| i.to_vec()), - event_tx, - kv_request_tx, - } - } -} - -impl ActorConfig { - /// Send a sleep intent - pub fn send_sleep_intent(&self) { - let event = protocol::make_actor_intent(rp::ActorIntent::ActorIntentSleep); - self.send_event(event); - } - - /// Send a stop intent - pub fn send_stop_intent(&self) { - let event = protocol::make_actor_intent(rp::ActorIntent::ActorIntentStop); - self.send_event(event); - } - - /// Set an alarm to wake at specified timestamp (milliseconds) - pub fn send_set_alarm(&self, alarm_ts: i64) { - let event = protocol::make_set_alarm(Some(alarm_ts)); - self.send_event(event); - } - - /// Clear the alarm - pub fn send_clear_alarm(&self) { - let event = protocol::make_set_alarm(None); - self.send_event(event); - } - - /// Send a custom event - fn send_event(&self, event: rp::Event) { - let actor_event = ActorEvent { - actor_id: self.actor_id.clone(), - generation: self.generation, - event, - }; - let _ = self.event_tx.send(actor_event); - } - - /// Send a KV get request - pub async fn send_kv_get(&self, keys: Vec>) -> Result { - let (response_tx, response_rx) = oneshot::channel(); - let request = KvRequest { - actor_id: self.actor_id.clone(), - data: rp::KvRequestData::KvGetRequest(rp::KvGetRequest { keys }), - response_tx, - }; - self.kv_request_tx - .send(request) - .map_err(|_| anyhow::anyhow!("failed to send KV get request"))?; - let response: rp::KvResponseData = response_rx - .await - .map_err(|_| anyhow::anyhow!("KV get request response channel closed"))?; - - match response { - rp::KvResponseData::KvGetResponse(data) => Ok(data), - rp::KvResponseData::KvErrorResponse(err) => { - Err(anyhow::anyhow!("KV get failed: {}", err.message)) - } - _ => Err(anyhow::anyhow!("unexpected response type for KV get")), - } - } - - /// Send a KV list request - pub async fn send_kv_list( - &self, - query: rp::KvListQuery, - reverse: Option, - limit: Option, - ) -> Result { - let (response_tx, response_rx) = oneshot::channel(); - let request = KvRequest { - actor_id: self.actor_id.clone(), - data: rp::KvRequestData::KvListRequest(rp::KvListRequest { - query, - reverse, - limit, - }), - response_tx, - }; - self.kv_request_tx - .send(request) - .map_err(|_| anyhow::anyhow!("failed to send KV list request"))?; - let response: rp::KvResponseData = response_rx - .await - .map_err(|_| anyhow::anyhow!("KV list request response channel closed"))?; - - match response { - rp::KvResponseData::KvListResponse(data) => Ok(data), - rp::KvResponseData::KvErrorResponse(err) => { - Err(anyhow::anyhow!("KV list failed: {}", err.message)) - } - _ => Err(anyhow::anyhow!("unexpected response type for KV list")), - } - } - - /// Send a KV put request - pub async fn send_kv_put(&self, keys: Vec>, values: Vec>) -> Result<()> { - let (response_tx, response_rx) = oneshot::channel(); - let request = KvRequest { - actor_id: self.actor_id.clone(), - data: rp::KvRequestData::KvPutRequest(rp::KvPutRequest { keys, values }), - response_tx, - }; - - self.kv_request_tx - .send(request) - .map_err(|_| anyhow::anyhow!("failed to send KV put request"))?; - - let response: rp::KvResponseData = response_rx - .await - .map_err(|_| anyhow::anyhow!("KV put request response channel closed"))?; - - match response { - rp::KvResponseData::KvPutResponse => Ok(()), - rp::KvResponseData::KvErrorResponse(err) => { - Err(anyhow::anyhow!("KV put failed: {}", err.message)) - } - _ => Err(anyhow::anyhow!("unexpected response type for KV put")), - } - } - - /// Send a KV delete request - pub async fn send_kv_delete(&self, keys: Vec>) -> Result<()> { - let (response_tx, response_rx) = oneshot::channel(); - let request = KvRequest { - actor_id: self.actor_id.clone(), - data: rp::KvRequestData::KvDeleteRequest(rp::KvDeleteRequest { keys }), - response_tx, - }; - self.kv_request_tx - .send(request) - .map_err(|_| anyhow::anyhow!("failed to send KV delete request"))?; - let response: rp::KvResponseData = response_rx - .await - .map_err(|_| anyhow::anyhow!("KV delete request response channel closed"))?; - - match response { - rp::KvResponseData::KvDeleteResponse => Ok(()), - rp::KvResponseData::KvErrorResponse(err) => { - Err(anyhow::anyhow!("KV delete failed: {}", err.message)) - } - _ => Err(anyhow::anyhow!("unexpected response type for KV delete")), - } - } - - /// Send a KV delete range request. - pub async fn send_kv_delete_range(&self, start: Vec, end: Vec) -> Result<()> { - let (response_tx, response_rx) = oneshot::channel(); - let request = KvRequest { - actor_id: self.actor_id.clone(), - data: rp::KvRequestData::KvDeleteRangeRequest(rp::KvDeleteRangeRequest { start, end }), - response_tx, - }; - self.kv_request_tx - .send(request) - .map_err(|_| anyhow::anyhow!("failed to send KV delete range request"))?; - let response: rp::KvResponseData = response_rx - .await - .map_err(|_| anyhow::anyhow!("KV delete range request response channel closed"))?; - - match response { - rp::KvResponseData::KvDeleteResponse => Ok(()), - rp::KvResponseData::KvErrorResponse(err) => { - Err(anyhow::anyhow!("KV delete range failed: {}", err.message)) - } - _ => Err(anyhow::anyhow!( - "unexpected response type for KV delete range" - )), - } - } - - /// Send a KV drop request - pub async fn send_kv_drop(&self) -> Result<()> { - let (response_tx, response_rx) = oneshot::channel(); - let request = KvRequest { - actor_id: self.actor_id.clone(), - data: rp::KvRequestData::KvDropRequest, - response_tx, - }; - self.kv_request_tx - .send(request) - .map_err(|_| anyhow::anyhow!("failed to send KV drop request"))?; - let response: rp::KvResponseData = response_rx - .await - .map_err(|_| anyhow::anyhow!("KV drop request response channel closed"))?; - - match response { - rp::KvResponseData::KvDropResponse => Ok(()), - rp::KvResponseData::KvErrorResponse(err) => { - Err(anyhow::anyhow!("KV drop failed: {}", err.message)) - } - _ => Err(anyhow::anyhow!("unexpected response type for KV drop")), - } - } -} - -/// Result of actor start operation -#[derive(Debug, Clone)] -pub enum ActorStartResult { - /// Send ActorStateRunning immediately - Running, - /// Wait specified duration before sending running - Delay(Duration), - /// Never send running (simulates timeout) - Timeout, - /// Crash immediately with exit code - Crash { code: i32, message: String }, -} - -/// Result of actor stop operation -#[derive(Debug, Clone)] -pub enum ActorStopResult { - /// Stop successfully (exit code 0) - Success, - /// Wait before stopping - Delay(Duration), - /// Crash with exit code - Crash { code: i32, message: String }, -} - -/// Trait for test actors that can be controlled programmatically -#[async_trait] -pub trait TestActor: Send + Sync { - /// Called when actor receives start command - async fn on_start(&mut self, config: ActorConfig) -> Result; - - /// Called when actor receives stop command - async fn on_stop(&mut self) -> Result; - - /// Called when actor receives alarm wake signal - async fn on_alarm(&mut self) -> Result<()> { - tracing::debug!("actor received alarm (default no-op)"); - Ok(()) - } - - /// Called when actor receives wake signal (from sleep) - async fn on_wake(&mut self) -> Result<()> { - tracing::debug!("actor received wake (default no-op)"); - Ok(()) - } - - /// Get actor's name for logging - fn name(&self) -> &str { - "TestActor" - } -} - -/// Events that actors can send directly via the event channel -#[derive(Debug, Clone)] -pub struct ActorEvent { - pub actor_id: String, - pub generation: u32, - pub event: rp::Event, -} - -/// KV requests that actors can send to the runner -pub struct KvRequest { - pub actor_id: String, - pub data: rp::KvRequestData, - pub response_tx: oneshot::Sender, -} diff --git a/engine/sdks/rust/engine-runner/src/behaviors.rs b/engine/sdks/rust/engine-runner/src/behaviors.rs deleted file mode 100644 index ee9973b7c7..0000000000 --- a/engine/sdks/rust/engine-runner/src/behaviors.rs +++ /dev/null @@ -1,582 +0,0 @@ -use crate::actor::*; -use anyhow::Result; -use async_trait::async_trait; -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -/// Simple echo actor that responds successfully and does nothing special -pub struct EchoActor; - -impl EchoActor { - pub fn new() -> Self { - Self {} - } -} - -impl Default for EchoActor { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl TestActor for EchoActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::info!(actor_id = ?config.actor_id, generation = config.generation, "echo actor started"); - Ok(ActorStartResult::Running) - } - - async fn on_stop(&mut self) -> Result { - tracing::info!("echo actor stopped"); - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "EchoActor" - } -} - -/// Actor that crashes immediately on start with specified exit code -pub struct CrashOnStartActor { - pub exit_code: i32, - pub message: String, - notify_tx: Option>>>>, -} - -impl CrashOnStartActor { - pub fn new(exit_code: i32) -> Self { - Self { - exit_code, - message: format!("crash on start with code {}", exit_code), - notify_tx: None, - } - } - - pub fn new_with_notify( - exit_code: i32, - notify_tx: std::sync::Arc>>>, - ) -> Self { - Self { - exit_code, - message: format!("crash on start with code {}", exit_code), - notify_tx: Some(notify_tx), - } - } -} - -#[async_trait] -impl TestActor for CrashOnStartActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::warn!( - actor_id = ?config.actor_id, - generation = config.generation, - exit_code = self.exit_code, - "crash on start actor crashing" - ); - - // Notify before crashing - if let Some(notify_tx) = &self.notify_tx { - let mut guard = notify_tx.lock().expect("failed to lock notify_tx"); - if let Some(tx) = guard.take() { - let _ = tx.send(()); - } - } - - Ok(ActorStartResult::Crash { - code: self.exit_code, - message: self.message.clone(), - }) - } - - async fn on_stop(&mut self) -> Result { - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "CrashOnStartActor" - } -} - -/// Actor that delays before sending running state -pub struct DelayedStartActor { - pub delay: Duration, -} - -impl DelayedStartActor { - pub fn new(delay: Duration) -> Self { - Self { delay } - } -} - -#[async_trait] -impl TestActor for DelayedStartActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::info!( - actor_id = ?config.actor_id, - generation = config.generation, - delay_ms = self.delay.as_millis(), - "delayed start actor will delay before running" - ); - Ok(ActorStartResult::Delay(self.delay)) - } - - async fn on_stop(&mut self) -> Result { - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "DelayedStartActor" - } -} - -/// Actor that never sends running state (simulates timeout) -pub struct TimeoutActor; - -impl TimeoutActor { - pub fn new() -> Self { - Self {} - } -} - -impl Default for TimeoutActor { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl TestActor for TimeoutActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::warn!( - actor_id = ?config.actor_id, - generation = config.generation, - "timeout actor will never send running state" - ); - Ok(ActorStartResult::Timeout) - } - - async fn on_stop(&mut self) -> Result { - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "TimeoutActor" - } -} - -/// Actor that sends sleep intent immediately after starting -pub struct SleepImmediatelyActor { - notify_tx: Option>>>>, -} - -impl SleepImmediatelyActor { - pub fn new() -> Self { - Self { notify_tx: None } - } - - pub fn new_with_notify( - notify_tx: std::sync::Arc>>>, - ) -> Self { - Self { - notify_tx: Some(notify_tx), - } - } -} - -impl Default for SleepImmediatelyActor { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl TestActor for SleepImmediatelyActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::info!( - actor_id = ?config.actor_id, - generation = config.generation, - "sleep immediately actor started, sending sleep intent" - ); - - // Send sleep intent immediately - config.send_sleep_intent(); - - // Notify that we're sending sleep intent - if let Some(notify_tx) = &self.notify_tx { - let mut guard = notify_tx.lock().expect("failed to lock notify_tx"); - if let Some(tx) = guard.take() { - let _ = tx.send(()); - } - } - - Ok(ActorStartResult::Running) - } - - async fn on_stop(&mut self) -> Result { - tracing::info!("sleep immediately actor stopped"); - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "SleepImmediatelyActor" - } -} - -/// Actor that sends stop intent immediately after starting -pub struct StopImmediatelyActor; - -impl StopImmediatelyActor { - pub fn new() -> Self { - Self - } -} - -impl Default for StopImmediatelyActor { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl TestActor for StopImmediatelyActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::info!( - actor_id = ?config.actor_id, - generation = config.generation, - "stop immediately actor started, sending stop intent" - ); - - // Send stop intent immediately - config.send_stop_intent(); - - Ok(ActorStartResult::Running) - } - - async fn on_stop(&mut self) -> Result { - tracing::info!("stop immediately actor stopped gracefully"); - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "StopImmediatelyActor" - } -} - -/// Actor that always crashes and increments a counter. -/// Used to test crash policy restart behavior. -pub struct CountingCrashActor { - crash_count: Arc, -} - -impl CountingCrashActor { - pub fn new(crash_count: Arc) -> Self { - Self { crash_count } - } -} - -#[async_trait] -impl TestActor for CountingCrashActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - let count = self - .crash_count - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - tracing::warn!( - actor_id = ?config.actor_id, - generation = config.generation, - crash_count = count + 1, - "counting crash actor crashing" - ); - Ok(ActorStartResult::Crash { - code: 1, - message: format!("crash #{}", count + 1), - }) - } - - async fn on_stop(&mut self) -> Result { - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "CountingCrashActor" - } -} - -/// Actor that crashes N times then succeeds -/// Used to test crash policy restart with retry reset on success -pub struct CrashNTimesThenSucceedActor { - crash_count: Arc>, - max_crashes: usize, -} - -impl CrashNTimesThenSucceedActor { - pub fn new(max_crashes: usize, crash_count: Arc>) -> Self { - Self { - crash_count, - max_crashes, - } - } -} - -#[async_trait] -impl TestActor for CrashNTimesThenSucceedActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - let mut count = self.crash_count.lock().unwrap(); - let current = *count; - - if current < self.max_crashes { - *count += 1; - tracing::warn!( - actor_id = ?config.actor_id, - generation = config.generation, - crash_count = current + 1, - max_crashes = self.max_crashes, - "crashing (will succeed after more crashes)" - ); - Ok(ActorStartResult::Crash { - code: 1, - message: format!("crash {} of {}", current + 1, self.max_crashes), - }) - } else { - tracing::info!( - actor_id = ?config.actor_id, - generation = config.generation, - crash_count = current, - "succeeded after crashes" - ); - Ok(ActorStartResult::Running) - } - } - - async fn on_stop(&mut self) -> Result { - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "CrashNTimesThenSucceedActor" - } -} - -/// Actor that notifies via a oneshot channel when it starts running -/// This allows tests to wait for the actor to actually start instead of sleeping -pub struct NotifyOnStartActor { - notify_tx: std::sync::Arc>>>, -} - -impl NotifyOnStartActor { - pub fn new( - notify_tx: std::sync::Arc>>>, - ) -> Self { - Self { notify_tx } - } -} - -#[async_trait] -impl TestActor for NotifyOnStartActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::info!( - actor_id = ?config.actor_id, - generation = config.generation, - "notify on start actor started, sending notification" - ); - - // Send notification that actor has started - let mut guard = self.notify_tx.lock().expect("failed to lock notify_tx"); - if let Some(tx) = guard.take() { - let _ = tx.send(()); - } - - Ok(ActorStartResult::Running) - } - - async fn on_stop(&mut self) -> Result { - tracing::info!("notify on start actor stopped"); - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "NotifyOnStartActor" - } -} - -/// Actor that verifies it received the expected input data -/// Crashes if input doesn't match or is missing, succeeds if it matches -pub struct VerifyInputActor { - expected_input: Vec, -} - -impl VerifyInputActor { - pub fn new(expected_input: Vec) -> Self { - Self { expected_input } - } -} - -#[async_trait] -impl TestActor for VerifyInputActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - tracing::info!( - actor_id = ?config.actor_id, - generation = config.generation, - expected_input_size = self.expected_input.len(), - received_input_size = config.input.as_ref().map(|i| i.len()), - "verify input actor started, checking input" - ); - - // Check if input is present - let Some(received_input) = &config.input else { - tracing::error!("no input data received"); - return Ok(ActorStartResult::Crash { - code: 1, - message: "no input data received".to_string(), - }); - }; - - // Check if input matches expected - if received_input != &self.expected_input { - tracing::error!( - expected_len = self.expected_input.len(), - received_len = received_input.len(), - "input data mismatch" - ); - return Ok(ActorStartResult::Crash { - code: 1, - message: format!( - "input mismatch: expected {} bytes, got {} bytes", - self.expected_input.len(), - received_input.len() - ), - }); - } - - tracing::info!("input data verified successfully"); - Ok(ActorStartResult::Running) - } - - async fn on_stop(&mut self) -> Result { - tracing::info!("verify input actor stopped"); - Ok(ActorStopResult::Success) - } - - fn name(&self) -> &str { - "VerifyInputActor" - } -} - -/// Generic actor that accepts closures for on_start and on_stop -/// This allows tests to define actor behavior inline without creating separate structs -pub struct CustomActor { - on_start_fn: Box< - dyn Fn( - ActorConfig, - ) -> std::pin::Pin< - Box> + Send>, - > + Send - + Sync, - >, - on_stop_fn: Box< - dyn Fn() -> std::pin::Pin< - Box> + Send>, - > + Send - + Sync, - >, -} - -/// Builder for CustomActor with default implementations -pub struct CustomActorBuilder { - on_start_fn: Option< - Box< - dyn Fn( - ActorConfig, - ) -> std::pin::Pin< - Box> + Send>, - > + Send - + Sync, - >, - >, - on_stop_fn: Option< - Box< - dyn Fn() -> std::pin::Pin< - Box> + Send>, - > + Send - + Sync, - >, - >, -} - -impl CustomActorBuilder { - pub fn new() -> Self { - Self { - on_start_fn: None, - on_stop_fn: None, - } - } - - pub fn on_start(mut self, f: F) -> Self - where - F: Fn( - ActorConfig, - ) -> std::pin::Pin< - Box> + Send>, - > + Send - + Sync - + 'static, - { - self.on_start_fn = Some(Box::new(f)); - self - } - - pub fn on_stop(mut self, f: F) -> Self - where - F: Fn() -> std::pin::Pin< - Box> + Send>, - > + Send - + Sync - + 'static, - { - self.on_stop_fn = Some(Box::new(f)); - self - } - - pub fn build(self) -> CustomActor { - CustomActor { - on_start_fn: self.on_start_fn.unwrap_or_else(|| { - Box::new(|_config| { - Box::pin(async { Ok(ActorStartResult::Running) }) - as std::pin::Pin< - Box> + Send>, - > - }) - }), - on_stop_fn: self.on_stop_fn.unwrap_or_else(|| { - Box::new(|| { - Box::pin(async { Ok(ActorStopResult::Success) }) - as std::pin::Pin< - Box> + Send>, - > - }) - }), - } - } -} - -impl Default for CustomActorBuilder { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl TestActor for CustomActor { - async fn on_start(&mut self, config: ActorConfig) -> Result { - (self.on_start_fn)(config).await - } - - async fn on_stop(&mut self) -> Result { - (self.on_stop_fn)().await - } - - fn name(&self) -> &str { - "CustomActor" - } -} diff --git a/engine/sdks/rust/engine-runner/src/lib.rs b/engine/sdks/rust/engine-runner/src/lib.rs deleted file mode 100644 index 8cd636fedb..0000000000 --- a/engine/sdks/rust/engine-runner/src/lib.rs +++ /dev/null @@ -1,46 +0,0 @@ -//! Rust-based engine runner for Rivet. -//! -//! This library provides a pure Rust implementation of a Rivet runner that can be fully controlled -//! programmatically, allowing simulation of: -//! - Actor crashes with specific exit codes -//! - Protocol timing issues (delays, timeouts) -//! - Custom protocol events (sleep, alarms, etc.) -//! - Runner disconnection/reconnection scenarios -//! -//! # Example -//! -//! ```ignore -//! use rivet_engine_runner::{Runner, RunnerConfig, EchoActor}; -//! -//! let config = RunnerConfig::builder() -//! .endpoint("http://127.0.0.1:8080") -//! .token("dev") -//! .namespace("my-namespace") -//! .runner_name("my-runner") -//! .runner_key("unique-key") -//! .build(); -//! -//! let mut runner = Runner::new(config)?; -//! runner.register_actor("echo", |_| Box::new(EchoActor::new())); -//! runner.start().await?; -//! ``` - -mod actor; -mod behaviors; -mod protocol; -mod runner; - -pub use actor::{ActorConfig, ActorEvent, ActorStartResult, ActorStopResult, KvRequest, TestActor}; -pub use behaviors::{ - CountingCrashActor, CrashNTimesThenSucceedActor, CrashOnStartActor, CustomActor, - CustomActorBuilder, DelayedStartActor, EchoActor, NotifyOnStartActor, SleepImmediatelyActor, - StopImmediatelyActor, TimeoutActor, VerifyInputActor, -}; -pub use protocol::PROTOCOL_VERSION; -pub use runner::{ - ActorLifecycleEvent, Runner, RunnerBuilder, RunnerBuilderLegacy, RunnerConfig, - RunnerConfigBuilder, -}; - -// Re-export commonly used types from the protocol -pub use rivet_runner_protocol::mk2 as protocol_types; diff --git a/engine/sdks/rust/engine-runner/src/protocol.rs b/engine/sdks/rust/engine-runner/src/protocol.rs deleted file mode 100644 index 958fb14ccd..0000000000 --- a/engine/sdks/rust/engine-runner/src/protocol.rs +++ /dev/null @@ -1,51 +0,0 @@ -use anyhow::Result; -use rivet_runner_protocol as rp; -use rivet_runner_protocol::mk2 as rp2; -use vbare::OwnedVersionedData; - -pub const PROTOCOL_VERSION: u16 = rp::PROTOCOL_MK2_VERSION; - -/// Helper to decode messages from server (MK2) -pub fn decode_to_client(buf: &[u8], protocol_version: u16) -> Result { - // Use versioned deserialization to handle protocol version properly - ::deserialize(buf, protocol_version) -} - -/// Helper to encode messages to server (MK2) -pub fn encode_to_server(msg: rp2::ToServer) -> Vec { - rp::versioned::ToServerMk2::wrap_latest(msg) - .serialize(PROTOCOL_VERSION) - .expect("failed to serialize ToServer") -} - -/// Helper to create event wrapper with checkpoint (MK2) -pub fn make_event_wrapper( - actor_id: &str, - generation: u32, - index: u64, - event: rp2::Event, -) -> rp2::EventWrapper { - rp2::EventWrapper { - checkpoint: rp2::ActorCheckpoint { - actor_id: actor_id.to_string(), - generation, - index: index as i64, - }, - inner: event, - } -} - -/// Helper to create actor state update event (MK2) -pub fn make_actor_state_update(state: rp2::ActorState) -> rp2::Event { - rp2::Event::EventActorStateUpdate(rp2::EventActorStateUpdate { state }) -} - -/// Helper to create actor intent event (MK2) -pub fn make_actor_intent(intent: rp2::ActorIntent) -> rp2::Event { - rp2::Event::EventActorIntent(rp2::EventActorIntent { intent }) -} - -/// Helper to create set alarm event (MK2) -pub fn make_set_alarm(alarm_ts: Option) -> rp2::Event { - rp2::Event::EventActorSetAlarm(rp2::EventActorSetAlarm { alarm_ts }) -} diff --git a/engine/sdks/rust/engine-runner/src/runner.rs b/engine/sdks/rust/engine-runner/src/runner.rs deleted file mode 100644 index 08340cf245..0000000000 --- a/engine/sdks/rust/engine-runner/src/runner.rs +++ /dev/null @@ -1,1079 +0,0 @@ -use crate::{actor::*, protocol}; -use anyhow::{Context, Result}; -use futures_util::{SinkExt, StreamExt}; -use rivet_runner_protocol::mk2 as rp; -use std::{ - collections::HashMap, - sync::{ - Arc, - atomic::{AtomicBool, Ordering}, - }, - time::Duration, -}; -use tokio::sync::{Mutex, broadcast, mpsc, oneshot}; -use tokio_tungstenite::{connect_async, tungstenite::Message}; - -const RUNNER_PING_INTERVAL: Duration = Duration::from_secs(15); - -type ActorFactory = Arc Box + Send + Sync>; -type WsStream = - tokio_tungstenite::WebSocketStream>; - -/// Lifecycle events for actors that tests can subscribe to -#[derive(Debug, Clone)] -pub enum ActorLifecycleEvent { - Started { actor_id: String, generation: u32 }, - Stopped { actor_id: String, generation: u32 }, -} - -/// Configuration for the engine runner. -/// -/// This matches the TypeScript RunnerConfig interface. -#[derive(Clone)] -pub struct RunnerConfig { - /// The endpoint URL to connect to (e.g., "http://127.0.0.1:8080") - pub endpoint: String, - /// Authentication token - pub token: String, - /// Namespace to connect to - pub namespace: String, - /// Name of this runner (machine-readable) - pub runner_name: String, - /// Unique key for this runner instance - pub runner_key: String, - /// Protocol version number - pub version: u32, - /// Total number of actor slots this runner supports - pub total_slots: u32, - /// Optional metadata to attach to the runner - pub metadata: Option, -} - -impl RunnerConfig { - /// Create a new builder for RunnerConfig - pub fn builder() -> RunnerConfigBuilder { - RunnerConfigBuilder::default() - } -} - -/// Builder for RunnerConfig -#[derive(Default)] -pub struct RunnerConfigBuilder { - endpoint: Option, - token: Option, - namespace: Option, - runner_name: Option, - runner_key: Option, - version: Option, - total_slots: Option, - metadata: Option, -} - -impl RunnerConfigBuilder { - pub fn endpoint(mut self, endpoint: impl Into) -> Self { - self.endpoint = Some(endpoint.into()); - self - } - - pub fn token(mut self, token: impl Into) -> Self { - self.token = Some(token.into()); - self - } - - pub fn namespace(mut self, namespace: impl Into) -> Self { - self.namespace = Some(namespace.into()); - self - } - - pub fn runner_name(mut self, name: impl Into) -> Self { - self.runner_name = Some(name.into()); - self - } - - pub fn runner_key(mut self, key: impl Into) -> Self { - self.runner_key = Some(key.into()); - self - } - - pub fn version(mut self, version: u32) -> Self { - self.version = Some(version); - self - } - - pub fn total_slots(mut self, slots: u32) -> Self { - self.total_slots = Some(slots); - self - } - - pub fn metadata(mut self, metadata: serde_json::Value) -> Self { - self.metadata = Some(metadata); - self - } - - pub fn build(self) -> Result { - Ok(RunnerConfig { - endpoint: self.endpoint.context("endpoint is required")?, - token: self.token.unwrap_or_else(|| "dev".to_string()), - namespace: self.namespace.context("namespace is required")?, - runner_name: self - .runner_name - .unwrap_or_else(|| "engine-runner".to_string()), - runner_key: self - .runner_key - .unwrap_or_else(|| format!("key-{:012x}", rand::random::())), - version: self.version.unwrap_or(1), - total_slots: self.total_slots.unwrap_or(100), - metadata: self.metadata, - }) - } -} - -/// Internal configuration with actor factories -#[derive(Clone)] -struct InternalConfig { - namespace: String, - runner_name: String, - runner_key: String, - version: u32, - total_slots: u32, - endpoint: String, - token: String, - actor_factories: HashMap, -} - -/// Engine runner for programmatic actor lifecycle control -pub struct Runner { - config: InternalConfig, - - // State - pub runner_id: Arc>>, - actors: Arc>>, - /// Per-actor event indices for MK2 checkpoints - actor_event_indices: Arc>>, - event_history: Arc>>, - shutdown: Arc, - is_child_task: bool, - - // Event channel for actors to push events - event_tx: mpsc::UnboundedSender, - event_rx: Arc>>, - - // KV request channel for actors to send KV requests - kv_request_tx: mpsc::UnboundedSender, - kv_request_rx: Arc>>, - next_kv_request_id: Arc>, - kv_pending_requests: Arc>>>, - - // Lifecycle event broadcast channel - lifecycle_tx: broadcast::Sender, - - // Shutdown channel - shutdown_tx: Arc>>>, -} - -struct ActorState { - #[allow(dead_code)] - actor_id: String, - #[allow(dead_code)] - generation: u32, - actor: Box, -} - -/// Builder for creating a Runner instance -pub struct RunnerBuilder { - config: RunnerConfig, - actor_factories: HashMap, -} - -impl RunnerBuilder { - /// Create a new RunnerBuilder with the given configuration - pub fn new(config: RunnerConfig) -> Self { - Self { - config, - actor_factories: HashMap::new(), - } - } - - /// Create a new RunnerBuilder from a namespace (for backwards compatibility) - /// - /// This is a convenience method that requires endpoint and token to be set later. - pub fn from_namespace(namespace: &str) -> RunnerBuilderLegacy { - RunnerBuilderLegacy { - namespace: namespace.to_string(), - runner_name: "engine-runner".to_string(), - runner_key: format!("key-{:012x}", rand::random::()), - version: 1, - total_slots: 100, - actor_factories: HashMap::new(), - } - } - - /// Register an actor factory for a specific actor name - pub fn with_actor_behavior(mut self, actor_name: &str, factory: F) -> Self - where - F: Fn(ActorConfig) -> Box + Send + Sync + 'static, - { - self.actor_factories - .insert(actor_name.to_string(), Arc::new(factory)); - self - } - - /// Build the Runner instance - pub fn build(self) -> Result { - let config = InternalConfig { - namespace: self.config.namespace, - runner_name: self.config.runner_name, - runner_key: self.config.runner_key, - version: self.config.version, - total_slots: self.config.total_slots, - endpoint: self.config.endpoint, - token: self.config.token, - actor_factories: self.actor_factories, - }; - - // Create event channel for actors to push events - let (event_tx, event_rx) = mpsc::unbounded_channel(); - - // Create KV request channel for actors to send KV requests - let (kv_request_tx, kv_request_rx) = mpsc::unbounded_channel(); - - // Create lifecycle event broadcast channel (capacity of 100 for buffering) - let (lifecycle_tx, _) = broadcast::channel(100); - - Ok(Runner { - config, - runner_id: Arc::new(Mutex::new(None)), - actors: Arc::new(Mutex::new(HashMap::new())), - actor_event_indices: Arc::new(Mutex::new(HashMap::new())), - event_history: Arc::new(Mutex::new(Vec::new())), - shutdown: Arc::new(AtomicBool::new(false)), - is_child_task: false, - event_tx, - event_rx: Arc::new(Mutex::new(event_rx)), - kv_request_tx, - kv_request_rx: Arc::new(Mutex::new(kv_request_rx)), - next_kv_request_id: Arc::new(Mutex::new(0)), - kv_pending_requests: Arc::new(Mutex::new(HashMap::new())), - lifecycle_tx, - shutdown_tx: Arc::new(Mutex::new(None)), - }) - } -} - -/// Legacy builder for backwards compatibility with test code -pub struct RunnerBuilderLegacy { - namespace: String, - runner_name: String, - runner_key: String, - version: u32, - total_slots: u32, - actor_factories: HashMap, -} - -impl RunnerBuilderLegacy { - pub fn with_runner_name(mut self, name: &str) -> Self { - self.runner_name = name.to_string(); - self - } - - pub fn with_runner_key(mut self, key: &str) -> Self { - self.runner_key = key.to_string(); - self - } - - pub fn with_version(mut self, version: u32) -> Self { - self.version = version; - self - } - - pub fn with_total_slots(mut self, total_slots: u32) -> Self { - self.total_slots = total_slots; - self - } - - /// Register an actor factory for a specific actor name - pub fn with_actor_behavior(mut self, actor_name: &str, factory: F) -> Self - where - F: Fn(ActorConfig) -> Box + Send + Sync + 'static, - { - self.actor_factories - .insert(actor_name.to_string(), Arc::new(factory)); - self - } - - /// Build the Runner with an endpoint and token - pub fn build_with_endpoint(self, endpoint: &str, token: &str) -> Result { - let config = InternalConfig { - namespace: self.namespace, - runner_name: self.runner_name, - runner_key: self.runner_key, - version: self.version, - total_slots: self.total_slots, - endpoint: endpoint.to_string(), - token: token.to_string(), - actor_factories: self.actor_factories, - }; - - // Create event channel for actors to push events - let (event_tx, event_rx) = mpsc::unbounded_channel(); - - // Create KV request channel for actors to send KV requests - let (kv_request_tx, kv_request_rx) = mpsc::unbounded_channel(); - - // Create lifecycle event broadcast channel (capacity of 100 for buffering) - let (lifecycle_tx, _) = broadcast::channel(100); - - Ok(Runner { - config, - runner_id: Arc::new(Mutex::new(None)), - actors: Arc::new(Mutex::new(HashMap::new())), - actor_event_indices: Arc::new(Mutex::new(HashMap::new())), - event_history: Arc::new(Mutex::new(Vec::new())), - shutdown: Arc::new(AtomicBool::new(false)), - is_child_task: false, - event_tx, - event_rx: Arc::new(Mutex::new(event_rx)), - kv_request_tx, - kv_request_rx: Arc::new(Mutex::new(kv_request_rx)), - next_kv_request_id: Arc::new(Mutex::new(0)), - kv_pending_requests: Arc::new(Mutex::new(HashMap::new())), - lifecycle_tx, - shutdown_tx: Arc::new(Mutex::new(None)), - }) - } -} - -impl Runner { - /// Subscribe to actor lifecycle events - pub fn subscribe_lifecycle_events(&self) -> broadcast::Receiver { - self.lifecycle_tx.subscribe() - } - - /// Start the runner - pub async fn start(&self) -> Result<()> { - tracing::info!( - namespace = %self.config.namespace, - runner_name = %self.config.runner_name, - runner_key = %self.config.runner_key, - "starting engine runner" - ); - - let ws_url = self.build_ws_url(); - - tracing::debug!(ws_url = %ws_url, "connecting to pegboard"); - - // Connect to WebSocket with protocols - let token_protocol = format!("rivet_token.{}", self.config.token); - - // Build the request properly with all WebSocket headers - use tokio_tungstenite::tungstenite::client::IntoClientRequest; - let mut request = ws_url - .into_client_request() - .context("failed to build WebSocket request")?; - - // Add the Sec-WebSocket-Protocol header - request.headers_mut().insert( - "Sec-WebSocket-Protocol", - format!("rivet, {}", token_protocol).parse().unwrap(), - ); - - let (ws_stream, _response) = connect_async(request) - .await - .context("failed to connect to WebSocket")?; - - tracing::info!("websocket connected"); - - // Create shutdown channel - let (shutdown_tx, shutdown_rx) = oneshot::channel(); - *self.shutdown_tx.lock().await = Some(shutdown_tx); - - // Clone self for the spawned task - let runner = self.clone_for_task(); - - tokio::spawn(async move { - if let Err(err) = runner.run_message_loop(ws_stream, shutdown_rx).await { - tracing::error!(?err, "engine runner message loop failed"); - } - }); - - Ok(()) - } - - /// Clone the runner for passing to async tasks - fn clone_for_task(&self) -> Self { - Self { - config: self.config.clone(), - runner_id: self.runner_id.clone(), - actors: self.actors.clone(), - actor_event_indices: self.actor_event_indices.clone(), - event_history: self.event_history.clone(), - is_child_task: true, - shutdown: self.shutdown.clone(), - event_tx: self.event_tx.clone(), - event_rx: self.event_rx.clone(), - kv_request_tx: self.kv_request_tx.clone(), - kv_request_rx: self.kv_request_rx.clone(), - next_kv_request_id: self.next_kv_request_id.clone(), - kv_pending_requests: self.kv_pending_requests.clone(), - lifecycle_tx: self.lifecycle_tx.clone(), - shutdown_tx: self.shutdown_tx.clone(), - } - } - - /// Wait for runner to be ready and return runner ID - pub async fn wait_ready(&self) -> String { - // Poll until runner_id is set - loop { - let runner_id = self.runner_id.lock().await; - if let Some(id) = runner_id.as_ref() { - // In MK2, we need to wait for the workflow to process the Init signal - // and mark the runner as eligible for actor allocation. - // This can take some time due to workflow processing: - // 1. Workflow receives Init signal - // 2. Workflow executes MarkEligible activity - // 3. Database is updated with runner allocation index - tokio::time::sleep(Duration::from_millis(2000)).await; - return id.clone(); - } - drop(runner_id); - tokio::time::sleep(Duration::from_millis(100)).await; - } - } - - /// Check if runner has an actor - pub async fn has_actor(&self, actor_id: &str) -> bool { - let actors = self.actors.lock().await; - actors.contains_key(actor_id) - } - - /// Get runner's current actor IDs - pub async fn get_actor_ids(&self) -> Vec { - let actors = self.actors.lock().await; - actors.keys().cloned().collect() - } - - pub fn name(&self) -> &str { - &self.config.runner_name - } - - /// Shutdown the runner gracefully (destroys actors first) - pub async fn shutdown(&self) { - tracing::info!("shutting down engine runner"); - self.shutdown.store(true, Ordering::SeqCst); - - // Send shutdown signal to close ws_stream - if let Some(tx) = self.shutdown_tx.lock().await.take() { - let _ = tx.send(()); - } - } - - /// Crash the runner without graceful shutdown. - /// This simulates an ungraceful disconnect where the runner stops responding - /// without destroying its actors first. Use this to test RunnerNoResponse errors. - pub async fn crash(&self) { - tracing::info!("crashing engine runner (ungraceful disconnect)"); - self.shutdown.store(true, Ordering::SeqCst); - - // Just drop the websocket without cleanup - don't send any signals - // The server will detect the disconnect and actors will remain in - // an unresponsive state until they timeout. - if let Some(tx) = self.shutdown_tx.lock().await.take() { - let _ = tx.send(()); - } - - // Clear local actor state without notifying server - self.actors.lock().await.clear(); - } - - fn build_ws_url(&self) -> String { - let ws_endpoint = self.config.endpoint.replace("http://", "ws://"); - format!( - "{}/runners/connect?protocol_version={}&namespace={}&runner_key={}", - ws_endpoint.trim_end_matches('/'), - protocol::PROTOCOL_VERSION, - urlencoding::encode(&self.config.namespace), - urlencoding::encode(&self.config.runner_key) - ) - } - - fn build_init_message(&self) -> rp::ToServer { - // MK2 init doesn't have lastCommandIdx - uses checkpoints instead - rp::ToServer::ToServerInit(rp::ToServerInit { - name: self.config.runner_name.clone(), - version: self.config.version, - total_slots: self.config.total_slots, - prepopulate_actor_names: None, - metadata: None, - }) - } - - async fn run_message_loop( - self, - mut ws_stream: WsStream, - mut shutdown_rx: oneshot::Receiver<()>, - ) -> Result<()> { - // Send init message - let init_msg = self.build_init_message(); - let encoded = protocol::encode_to_server(init_msg); - ws_stream - .send(Message::Binary(encoded.into())) - .await - .context("failed to send init message")?; - - tracing::debug!("sent init message"); - - let mut ping_interval = tokio::time::interval(RUNNER_PING_INTERVAL); - // We lock here as these rx's are only for run_message_loop - let mut event_rx = self.event_rx.lock().await; - let mut kv_request_rx = self.kv_request_rx.lock().await; - - loop { - tokio::select! { - biased; - _ = &mut shutdown_rx => { - tracing::info!("received shutdown signal, closing websocket"); - let _ = ws_stream.close(None).await; - break; - } - - _ = ping_interval.tick() => { - if self.shutdown.load(Ordering::SeqCst) { - break; - } - - // Send pong (MK2 uses ToServerPong instead of ToServerPing) - let pong = rp::ToServer::ToServerPong(rp::ToServerPong { - ts: chrono::Utc::now().timestamp_millis(), - }); - let encoded = protocol::encode_to_server(pong); - ws_stream.send(Message::Binary(encoded.into())).await?; - } - - // Listen for events pushed from actors - Some(actor_event) = event_rx.recv() => { - if self.shutdown.load(Ordering::SeqCst) { - tracing::info!("shutting down"); - break; - } - - tracing::debug!( - actor_id = ?actor_event.actor_id, - generation = actor_event.generation, - "received event from actor" - ); - - self.send_actor_event(&mut ws_stream, actor_event).await?; - } - - // Listen for KV requests from actors - Some(kv_request) = kv_request_rx.recv() => { - if self.shutdown.load(Ordering::SeqCst) { - break; - } - - tracing::debug!( - actor_id = ?kv_request.actor_id, - "received kv request from actor" - ); - - self.send_kv_request(&mut ws_stream, kv_request).await?; - } - - msg = ws_stream.next() => { - if self.shutdown.load(Ordering::SeqCst) { - break; - } - - match msg { - Some(std::result::Result::Ok(Message::Binary(buf))) => { - self.handle_message(&mut ws_stream, &buf).await?; - } - Some(std::result::Result::Ok(Message::Close(_))) => { - tracing::info!("websocket closed by server"); - break; - } - Some(std::result::Result::Err(err)) => { - tracing::error!(?err, "websocket error"); - return Err(err.into()); - } - None => { - tracing::info!("websocket stream ended"); - break; - } - _ => {} - } - } - } - } - - tracing::info!("engine runner message loop exiting"); - Ok(()) - } - - /// Send an event pushed from an actor - async fn send_actor_event( - &self, - ws_stream: &mut WsStream, - actor_event: ActorEvent, - ) -> Result<()> { - // Get next event index for this actor (MK2 uses per-actor checkpoints) - let mut indices = self.actor_event_indices.lock().await; - let idx = indices.entry(actor_event.actor_id.clone()).or_insert(-1); - *idx += 1; - let event_idx = *idx; - drop(indices); - - let event_wrapper = protocol::make_event_wrapper( - &actor_event.actor_id, - actor_event.generation, - event_idx as u64, - actor_event.event, - ); - - self.event_history.lock().await.push(event_wrapper.clone()); - - tracing::debug!( - actor_id = ?actor_event.actor_id, - generation = actor_event.generation, - event_idx = event_idx, - "sending actor event" - ); - - let msg = rp::ToServer::ToServerEvents(vec![event_wrapper]); - let encoded = protocol::encode_to_server(msg); - ws_stream.send(Message::Binary(encoded.into())).await?; - - Ok(()) - } - - async fn handle_message(&self, ws_stream: &mut WsStream, buf: &[u8]) -> Result<()> { - let msg = protocol::decode_to_client(buf, protocol::PROTOCOL_VERSION)?; - - match msg { - rp::ToClient::ToClientInit(init) => { - self.handle_init(init, ws_stream).await?; - } - rp::ToClient::ToClientCommands(commands) => { - self.handle_commands(commands, ws_stream).await?; - } - rp::ToClient::ToClientAckEvents(ack) => { - self.handle_ack_events(ack).await; - } - rp::ToClient::ToClientKvResponse(response) => { - self.handle_kv_response(response).await; - } - _ => { - tracing::debug!(?msg, "ignoring message type"); - } - } - - Ok(()) - } - - async fn handle_init(&self, init: rp::ToClientInit, _ws_stream: &mut WsStream) -> Result<()> { - tracing::info!( - runner_id = %init.runner_id, - "received init from server" - ); - - *self.runner_id.lock().await = Some(init.runner_id.clone()); - - // MK2 doesn't have lastEventIdx in init - events are acked via checkpoints - // For simplicity, we don't resend events on reconnect in the engine runner - - Ok(()) - } - - async fn handle_commands( - &self, - commands: Vec, - ws_stream: &mut WsStream, - ) -> Result<()> { - tracing::info!(count = commands.len(), "received commands"); - - for cmd_wrapper in commands { - let checkpoint = &cmd_wrapper.checkpoint; - tracing::debug!( - actor_id = %checkpoint.actor_id, - generation = checkpoint.generation, - index = checkpoint.index, - command = ?cmd_wrapper.inner, - "processing command" - ); - - match cmd_wrapper.inner { - rp::Command::CommandStartActor(start_cmd) => { - self.handle_start_actor( - checkpoint.actor_id.clone(), - checkpoint.generation, - start_cmd, - ws_stream, - ) - .await?; - } - rp::Command::CommandStopActor => { - // MK2 CommandStopActor is void - actor info is in checkpoint - self.handle_stop_actor( - checkpoint.actor_id.clone(), - checkpoint.generation, - ws_stream, - ) - .await?; - } - } - } - - Ok(()) - } - - async fn handle_start_actor( - &self, - actor_id: String, - generation: u32, - cmd: rp::CommandStartActor, - _ws_stream: &mut WsStream, - ) -> Result<()> { - tracing::info!(?actor_id, generation, name = %cmd.config.name, "starting actor"); - - // Create actor config - let config = ActorConfig::new( - &cmd.config, - actor_id.clone(), - generation, - self.event_tx.clone(), - self.kv_request_tx.clone(), - ); - - // Get factory for this actor name - let factory = self - .config - .actor_factories - .get(&cmd.config.name) - .context(format!( - "no factory registered for actor name: {}", - cmd.config.name - ))? - .clone(); - - // Clone self for the spawned task - let runner = self.clone_for_task(); - let actor_id_clone = actor_id.clone(); - - // Spawn actor execution in separate task to avoid blocking message loop - tokio::spawn(async move { - // Create actor - let mut actor = factory(config.clone()); - - tracing::debug!( - ?actor_id, - generation, - actor_type = actor.name(), - "created actor instance" - ); - - // Call on_start - let start_result = match actor.on_start(config).await { - std::result::Result::Ok(result) => result, - Err(err) => { - tracing::error!(?actor_id_clone, generation, ?err, "actor on_start failed"); - return; - } - }; - - tracing::debug!( - ?actor_id_clone, - generation, - ?start_result, - "actor on_start completed" - ); - - runner - .handle_actor_start_result(actor_id_clone, generation, actor, start_result) - .await; - }); - - Ok(()) - } - - async fn handle_actor_start_result( - &self, - actor_id: String, - generation: u32, - actor: Box, - start_result: ActorStartResult, - ) { - // Broadcast lifecycle event - tracing::info!("lifecycle_tx start"); - let _ = self.lifecycle_tx.send(ActorLifecycleEvent::Started { - actor_id: actor_id.clone(), - generation, - }); - - // Store actor - let actor_state = ActorState { - actor_id: actor_id.clone(), - generation, - actor, - }; - self.actors - .lock() - .await - .insert(actor_id.clone(), actor_state); - - // Handle start result and send state update via event - match start_result { - ActorStartResult::Running => { - let event = protocol::make_actor_state_update(rp::ActorState::ActorStateRunning); - self.event_tx - .send(ActorEvent { - actor_id: actor_id.clone(), - generation, - event, - }) - .expect("failed to send state update"); - } - ActorStartResult::Delay(duration) => { - let actor_id_clone = actor_id.clone(); - let event_tx = self.event_tx.clone(); - tokio::spawn(async move { - tracing::info!( - ?actor_id_clone, - generation, - delay_ms = duration.as_millis(), - "delaying before sending running state" - ); - tokio::time::sleep(duration).await; - let event = - protocol::make_actor_state_update(rp::ActorState::ActorStateRunning); - event_tx - .send(ActorEvent { - actor_id: actor_id_clone, - generation, - event, - }) - .expect("failed to send delayed state update"); - }); - } - ActorStartResult::Timeout => { - tracing::warn!( - ?actor_id, - generation, - "actor will timeout (not sending running)" - ); - // Don't send running state - } - ActorStartResult::Crash { code, message } => { - tracing::warn!(?actor_id, generation, code, %message, "actor crashed on start"); - let event = protocol::make_actor_state_update(rp::ActorState::ActorStateStopped( - rp::ActorStateStopped { - code: if code == 0 { - rp::StopCode::Ok - } else { - rp::StopCode::Error - }, - message: Some(message), - }, - )); - let _ = self - .event_tx - .send(ActorEvent { - actor_id: actor_id.clone(), - generation, - event, - }) - .expect("failed to send crash state update"); - - // Remove actor - self.actors.lock().await.remove(&actor_id); - } - } - } - - async fn handle_stop_actor( - &self, - actor_id: String, - generation: u32, - ws_stream: &mut WsStream, - ) -> Result<()> { - tracing::info!(?actor_id, generation, "stopping actor"); - - // Get actor - let mut actors_guard = self.actors.lock().await; - let actor_state = actors_guard.get_mut(&actor_id).context("actor not found")?; - - // Call on_stop - let stop_result = actor_state - .actor - .on_stop() - .await - .context("actor on_stop failed")?; - - tracing::debug!( - ?actor_id, - generation, - ?stop_result, - "actor on_stop completed" - ); - - // Broadcast lifecycle event - let _ = self.lifecycle_tx.send(ActorLifecycleEvent::Stopped { - actor_id: actor_id.clone(), - generation, - }); - - // Handle stop result - match stop_result { - ActorStopResult::Success => { - self.send_actor_state_update( - &actor_id, - generation, - rp::ActorState::ActorStateStopped(rp::ActorStateStopped { - code: rp::StopCode::Ok, - message: None, - }), - ws_stream, - ) - .await?; - } - ActorStopResult::Delay(duration) => { - tracing::info!(?actor_id, generation, ?duration, "delaying stop"); - tokio::time::sleep(duration).await; - self.send_actor_state_update( - &actor_id, - generation, - rp::ActorState::ActorStateStopped(rp::ActorStateStopped { - code: rp::StopCode::Ok, - message: None, - }), - ws_stream, - ) - .await?; - } - ActorStopResult::Crash { code, message } => { - tracing::warn!(?actor_id, generation, code, %message, "actor crashed on stop"); - self.send_actor_state_update( - &actor_id, - generation, - rp::ActorState::ActorStateStopped(rp::ActorStateStopped { - code: if code == 0 { - rp::StopCode::Ok - } else { - rp::StopCode::Error - }, - message: Some(message), - }), - ws_stream, - ) - .await?; - } - } - - // Remove actor - actors_guard.remove(&actor_id); - - Ok(()) - } - - async fn handle_ack_events(&self, ack: rp::ToClientAckEvents) { - // MK2 uses per-actor checkpoints for acknowledgments - let checkpoints = &ack.last_event_checkpoints; - - let mut events = self.event_history.lock().await; - let original_len = events.len(); - - // Remove events that have been acknowledged based on checkpoints - events.retain(|e| { - // Check if this event's checkpoint is covered by any ack checkpoint - !checkpoints.iter().any(|ck| { - ck.actor_id == e.checkpoint.actor_id - && ck.generation == e.checkpoint.generation - && ck.index >= e.checkpoint.index - }) - }); - - let pruned = original_len - events.len(); - if pruned > 0 { - tracing::debug!( - checkpoint_count = checkpoints.len(), - pruned, - "pruned acknowledged events" - ); - } - } - - async fn send_actor_state_update( - &self, - actor_id: &str, - generation: u32, - state: rp::ActorState, - ws_stream: &mut WsStream, - ) -> Result<()> { - let event = protocol::make_actor_state_update(state); - - self.send_actor_event( - ws_stream, - ActorEvent { - actor_id: actor_id.to_string(), - generation, - event, - }, - ) - .await?; - - Ok(()) - } - - async fn send_kv_request(&self, ws_stream: &mut WsStream, kv_request: KvRequest) -> Result<()> { - let mut request_id = self.next_kv_request_id.lock().await; - let id = *request_id; - *request_id += 1; - drop(request_id); - - // Store the response channel - self.kv_pending_requests - .lock() - .await - .insert(id, kv_request.response_tx); - - tracing::debug!( - actor_id = ?kv_request.actor_id, - request_id = id, - "sending kv request" - ); - - let msg = rp::ToServer::ToServerKvRequest(rp::ToServerKvRequest { - actor_id: kv_request.actor_id, - request_id: id, - data: kv_request.data, - }); - let encoded = protocol::encode_to_server(msg); - ws_stream.send(Message::Binary(encoded.into())).await?; - - Ok(()) - } - - async fn handle_kv_response(&self, response: rp::ToClientKvResponse) { - let request_id = response.request_id; - - tracing::debug!(request_id, "received kv response"); - - let response_tx = self.kv_pending_requests.lock().await.remove(&request_id); - - if let Some(tx) = response_tx { - let _ = tx.send(response.data); - } else { - tracing::warn!(request_id, "received kv response for unknown request id"); - } - } -} - -impl Drop for Runner { - fn drop(&mut self) { - if self.is_child_task { - return; - } - // Signal shutdown when runner is dropped - self.shutdown.store(true, Ordering::SeqCst); - tracing::debug!("engine runner dropped, shutdown signaled"); - } -} diff --git a/engine/sdks/rust/envoy-client/README.md b/engine/sdks/rust/envoy-client/README.md deleted file mode 100644 index 7b311d41bd..0000000000 --- a/engine/sdks/rust/envoy-client/README.md +++ /dev/null @@ -1 +0,0 @@ -TODO: Rewrite this entire piece of garbage without ai \ No newline at end of file diff --git a/engine/sdks/rust/envoy-client/src/lib.rs b/engine/sdks/rust/envoy-client/src/lib.rs deleted file mode 100644 index f765806d3e..0000000000 --- a/engine/sdks/rust/envoy-client/src/lib.rs +++ /dev/null @@ -1,41 +0,0 @@ -//! Rust-based envoy client for Rivet. -//! -//! This library provides a pure Rust implementation of a Rivet envoy that can be fully controlled -//! programmatically, allowing simulation of: -//! - Actor crashes with specific exit codes -//! - Protocol timing issues (delays, timeouts) -//! - Custom protocol events (sleep, alarms, etc.) -//! - Envoy disconnection/reconnection scenarios -//! -//! # Example -//! -//! ```ignore -//! use rivet_envoy_client::{Envoy, EnvoyConfig, EchoActor}; -//! -//! let config = EnvoyConfig::builder() -//! .endpoint("http://127.0.0.1:8080") -//! .token("dev") -//! .namespace("my-namespace") -//! .pool_name("my-pool") -//! .build(); -//! -//! let mut envoy = Envoy::new(config)?; -//! envoy.register_actor("echo", |_| Box::new(EchoActor::new())); -//! envoy.start().await?; -//! ``` - -mod actor; -mod behaviors; -mod envoy; -mod utils; - -pub use actor::{ActorConfig, ActorEvent, ActorStartResult, ActorStopResult, KvRequest, TestActor}; -pub use behaviors::{ - CountingCrashActor, CrashNTimesThenSucceedActor, CrashOnStartActor, CustomActor, - CustomActorBuilder, DelayedStartActor, EchoActor, NotifyOnStartActor, SleepImmediatelyActor, - StopImmediatelyActor, TimeoutActor, VerifyInputActor, -}; -pub use envoy::{ActorLifecycleEvent, Envoy, EnvoyBuilder, EnvoyConfig, EnvoyConfigBuilder}; - -// Re-export commonly used types from the protocol -pub use rivet_envoy_protocol as protocol; diff --git a/engine/sdks/rust/envoy-client/Cargo.toml b/engine/sdks/rust/test-envoy/Cargo.toml similarity index 54% rename from engine/sdks/rust/envoy-client/Cargo.toml rename to engine/sdks/rust/test-envoy/Cargo.toml index ef4e865eda..89d9b61d27 100644 --- a/engine/sdks/rust/envoy-client/Cargo.toml +++ b/engine/sdks/rust/test-envoy/Cargo.toml @@ -1,23 +1,34 @@ [package] -name = "rivet-envoy-client" +name = "rivet-test-envoy" version.workspace = true authors.workspace = true license.workspace = true edition.workspace = true -description = "Rust-based envoy client for Rivet, enabling programmatic actor lifecycle control" +description = "Rust test envoy process and harness for pegboard actor testing" + +[[bin]] +name = "rivet-test-envoy" +path = "src/main.rs" [dependencies] anyhow.workspace = true +async-stream.workspace = true async-trait.workspace = true +axum.workspace = true chrono.workspace = true futures-util.workspace = true -rand.workspace = true +hex.workspace = true +reqwest.workspace = true rivet-envoy-protocol.workspace = true +rivet-runner-protocol.workspace = true +rivet-util.workspace = true +serde.workspace = true serde_bare.workspace = true serde_json.workspace = true -serde.workspace = true +tokio-stream.workspace = true tokio-tungstenite.workspace = true tokio.workspace = true +tracing-subscriber = { workspace = true, features = ["env-filter"] } tracing.workspace = true urlencoding.workspace = true uuid.workspace = true diff --git a/engine/sdks/rust/test-envoy/Dockerfile b/engine/sdks/rust/test-envoy/Dockerfile new file mode 100644 index 0000000000..8ee79f1efc --- /dev/null +++ b/engine/sdks/rust/test-envoy/Dockerfile @@ -0,0 +1,22 @@ +FROM rust:1.88-slim AS builder + +WORKDIR /app + +COPY Cargo.toml Cargo.lock ./ +COPY engine ./engine + +RUN cargo build --release -p rivet-test-envoy + +FROM debian:bookworm-slim + +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates curl \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +COPY --from=builder /app/target/release/rivet-test-envoy /usr/local/bin/rivet-test-envoy + +EXPOSE 5050 + +CMD ["rivet-test-envoy"] diff --git a/engine/sdks/rust/envoy-client/src/actor.rs b/engine/sdks/rust/test-envoy/src/actor.rs similarity index 88% rename from engine/sdks/rust/envoy-client/src/actor.rs rename to engine/sdks/rust/test-envoy/src/actor.rs index 9d7736785f..55d0f5c549 100644 --- a/engine/sdks/rust/envoy-client/src/actor.rs +++ b/engine/sdks/rust/test-envoy/src/actor.rs @@ -1,6 +1,7 @@ use anyhow::Result; use async_trait::async_trait; use rivet_envoy_protocol as protocol; +use rivet_runner_protocol::mk2 as runner_protocol; use std::time::Duration; use tokio::sync::{mpsc, oneshot}; @@ -45,6 +46,11 @@ impl ActorConfig { } impl ActorConfig { + /// Converts compatible KV list queries into the envoy protocol shape. + fn convert_kv_list_query(query: impl IntoEnvoyKvListQuery) -> protocol::KvListQuery { + query.into_envoy_kv_list_query() + } + /// Send a sleep intent pub fn send_sleep_intent(&self) { let event = utils::make_actor_intent(protocol::ActorIntent::ActorIntentSleep); @@ -106,7 +112,7 @@ impl ActorConfig { /// Send a KV list request pub async fn send_kv_list( &self, - query: protocol::KvListQuery, + query: impl IntoEnvoyKvListQuery, reverse: Option, limit: Option, ) -> Result { @@ -114,7 +120,7 @@ impl ActorConfig { let request = KvRequest { actor_id: self.actor_id.clone(), data: protocol::KvRequestData::KvListRequest(protocol::KvListRequest { - query, + query: Self::convert_kv_list_query(query), reverse, limit, }), @@ -240,6 +246,36 @@ impl ActorConfig { } } +pub trait IntoEnvoyKvListQuery { + fn into_envoy_kv_list_query(self) -> protocol::KvListQuery; +} + +impl IntoEnvoyKvListQuery for protocol::KvListQuery { + fn into_envoy_kv_list_query(self) -> protocol::KvListQuery { + self + } +} + +impl IntoEnvoyKvListQuery for runner_protocol::KvListQuery { + fn into_envoy_kv_list_query(self) -> protocol::KvListQuery { + match self { + runner_protocol::KvListQuery::KvListAllQuery => protocol::KvListQuery::KvListAllQuery, + runner_protocol::KvListQuery::KvListPrefixQuery(prefix) => { + protocol::KvListQuery::KvListPrefixQuery(protocol::KvListPrefixQuery { + key: prefix.key, + }) + } + runner_protocol::KvListQuery::KvListRangeQuery(range) => { + protocol::KvListQuery::KvListRangeQuery(protocol::KvListRangeQuery { + start: range.start, + end: range.end, + exclusive: range.exclusive, + }) + } + } + } +} + /// Result of actor start operation #[derive(Debug, Clone)] pub enum ActorStartResult { diff --git a/engine/sdks/rust/envoy-client/src/behaviors.rs b/engine/sdks/rust/test-envoy/src/behaviors.rs similarity index 100% rename from engine/sdks/rust/envoy-client/src/behaviors.rs rename to engine/sdks/rust/test-envoy/src/behaviors.rs diff --git a/engine/sdks/rust/envoy-client/src/envoy.rs b/engine/sdks/rust/test-envoy/src/envoy.rs similarity index 66% rename from engine/sdks/rust/envoy-client/src/envoy.rs rename to engine/sdks/rust/test-envoy/src/envoy.rs index ed05b69ebe..46b4568859 100644 --- a/engine/sdks/rust/envoy-client/src/envoy.rs +++ b/engine/sdks/rust/test-envoy/src/envoy.rs @@ -1,6 +1,7 @@ use anyhow::{Context, Result}; use futures_util::{SinkExt, StreamExt}; use rivet_envoy_protocol::{self as protocol, PROTOCOL_VERSION}; +use rivet_util::serde::HashableMap; use std::{ collections::HashMap, sync::{ @@ -19,6 +20,30 @@ type ActorFactory = Arc Box + Send + Sync> type WsStream = tokio_tungstenite::WebSocketStream>; +#[derive(Debug, Clone)] +struct HttpRequestState { + actor_id: String, + _method: String, + path: String, + _headers: HashableMap, + body: Vec, + stream: bool, +} + +#[derive(Debug, Clone)] +enum TunnelRequestKind { + Http(HttpRequestState), + WebSocket, +} + +#[derive(Debug, Clone)] +struct TunnelRequestState { + gateway_id: protocol::GatewayId, + request_id: protocol::RequestId, + next_message_index: u16, + kind: TunnelRequestKind, +} + /// Lifecycle events for actors that tests can subscribe to #[derive(Debug, Clone)] pub enum ActorLifecycleEvent { @@ -115,6 +140,7 @@ struct InternalConfig { endpoint: String, token: String, actor_factories: HashMap, + default_actor_factory: Option, } /// Envoy client for programmatic actor lifecycle control @@ -144,6 +170,9 @@ pub struct Envoy { // Lifecycle event broadcast channel lifecycle_tx: broadcast::Sender, + // HTTP/WebSocket tunnel state + tunnel_requests: Arc>>, + // Shutdown channel shutdown_tx: Arc>>>, } @@ -160,6 +189,7 @@ struct ActorState { pub struct EnvoyBuilder { config: EnvoyConfig, actor_factories: HashMap, + default_actor_factory: Option, } impl EnvoyBuilder { @@ -168,6 +198,7 @@ impl EnvoyBuilder { Self { config, actor_factories: HashMap::new(), + default_actor_factory: None, } } @@ -181,6 +212,15 @@ impl EnvoyBuilder { self } + /// Register a fallback actor factory used when no actor-specific behavior is registered. + pub fn with_default_actor_behavior(mut self, factory: F) -> Self + where + F: Fn(ActorConfig) -> Box + Send + Sync + 'static, + { + self.default_actor_factory = Some(Arc::new(factory)); + self + } + /// Build the Envoy instance pub fn build(self) -> Result { let config = InternalConfig { @@ -190,6 +230,7 @@ impl EnvoyBuilder { endpoint: self.config.endpoint, token: self.config.token, actor_factories: self.actor_factories, + default_actor_factory: self.default_actor_factory, }; // Create event channel for actors to push events @@ -217,6 +258,7 @@ impl EnvoyBuilder { next_kv_request_id: Arc::new(Mutex::new(0)), kv_pending_requests: Arc::new(Mutex::new(HashMap::new())), lifecycle_tx, + tunnel_requests: Arc::new(Mutex::new(HashMap::new())), shutdown_tx: Arc::new(Mutex::new(None)), }) } @@ -296,6 +338,7 @@ impl Envoy { next_kv_request_id: self.next_kv_request_id.clone(), kv_pending_requests: self.kv_pending_requests.clone(), lifecycle_tx: self.lifecycle_tx.clone(), + tunnel_requests: self.tunnel_requests.clone(), shutdown_tx: self.shutdown_tx.clone(), } } @@ -326,6 +369,60 @@ impl Envoy { &self.config.pool_name } + /// Send a sleep intent for the latest generation of an actor. + pub async fn sleep_actor(&self, actor_id: &str) { + let generation = { + let actors = self.actors.lock().await; + actors.get(actor_id).map(|actor| actor.generation) + }; + + if let Some(generation) = generation { + let _ = self.event_tx.send(ActorEvent { + actor_id: actor_id.to_string(), + generation, + event: utils::make_actor_intent(protocol::ActorIntent::ActorIntentSleep), + }); + } + } + + /// Start a serverless actor from the payload passed to `/api/rivet/start`. + pub async fn start_serverless_actor(&self, payload: &[u8]) -> Result<()> { + if payload.len() < 2 { + anyhow::bail!("serverless payload missing protocol version"); + } + + let version = u16::from_le_bytes([payload[0], payload[1]]); + if version != PROTOCOL_VERSION { + anyhow::bail!( + "serverless payload protocol version mismatch: {} vs {}", + version, + PROTOCOL_VERSION + ); + } + + let message = utils::decode_to_envoy(&payload[2..], version)?; + let protocol::ToEnvoy::ToEnvoyCommands(commands) = message else { + anyhow::bail!("invalid serverless payload"); + }; + if commands.len() != 1 { + anyhow::bail!("invalid serverless payload"); + } + + let command = commands.into_iter().next().expect("checked single command"); + let checkpoint = command.checkpoint; + let protocol::Command::CommandStartActor(start_cmd) = command.inner else { + anyhow::bail!("invalid serverless payload"); + }; + + self + .handle_start_actor( + checkpoint.actor_id, + checkpoint.generation, + start_cmd, + ) + .await + } + /// Shutdown the envoy gracefully (destroys actors first) pub async fn shutdown(&self) { tracing::info!("shutting down envoy client"); @@ -520,8 +617,8 @@ impl Envoy { protocol::ToEnvoy::ToEnvoyPing(ping) => { self.handle_ping(ws_stream, ping).await?; } - _ => { - tracing::debug!(?msg, "ignoring message type"); + protocol::ToEnvoy::ToEnvoyTunnelMessage(message) => { + self.handle_tunnel_message(ws_stream, message).await?; } } @@ -542,11 +639,7 @@ impl Envoy { Ok(()) } - async fn handle_commands( - &self, - commands: Vec, - ws_stream: &mut WsStream, - ) -> Result<()> { + async fn handle_commands(&self, commands: Vec, ws_stream: &mut WsStream) -> Result<()> { tracing::info!(count = commands.len(), "received commands"); for cmd_wrapper in commands { @@ -565,7 +658,6 @@ impl Envoy { checkpoint.actor_id.clone(), checkpoint.generation, start_cmd, - ws_stream, ) .await?; } @@ -588,7 +680,6 @@ impl Envoy { actor_id: String, generation: u32, cmd: protocol::CommandStartActor, - _ws_stream: &mut WsStream, ) -> Result<()> { tracing::info!(?actor_id, generation, name = %cmd.config.name, "starting actor"); @@ -606,11 +697,12 @@ impl Envoy { .config .actor_factories .get(&cmd.config.name) + .cloned() + .or_else(|| self.config.default_actor_factory.clone()) .context(format!( "no factory registered for actor name: {}", cmd.config.name - ))? - .clone(); + ))?; // Clone self for the spawned task let envoy = self.clone_for_task(); @@ -869,6 +961,378 @@ impl Envoy { Ok(()) } + async fn handle_tunnel_message( + &self, + ws_stream: &mut WsStream, + message: protocol::ToEnvoyTunnelMessage, + ) -> Result<()> { + match message.message_kind { + protocol::ToEnvoyTunnelMessageKind::ToEnvoyRequestStart(req) => { + self + .handle_request_start(ws_stream, message.message_id, req) + .await?; + } + protocol::ToEnvoyTunnelMessageKind::ToEnvoyRequestChunk(chunk) => { + self + .handle_request_chunk(ws_stream, message.message_id, chunk) + .await?; + } + protocol::ToEnvoyTunnelMessageKind::ToEnvoyRequestAbort => { + self.handle_request_abort(message.message_id).await; + } + protocol::ToEnvoyTunnelMessageKind::ToEnvoyWebSocketOpen(open) => { + self + .handle_websocket_open(ws_stream, message.message_id, open) + .await?; + } + protocol::ToEnvoyTunnelMessageKind::ToEnvoyWebSocketMessage(msg) => { + self + .handle_websocket_message(ws_stream, message.message_id, msg) + .await?; + } + protocol::ToEnvoyTunnelMessageKind::ToEnvoyWebSocketClose(close) => { + self + .handle_websocket_close(ws_stream, message.message_id, close) + .await?; + } + } + + Ok(()) + } + + async fn handle_request_start( + &self, + ws_stream: &mut WsStream, + message_id: protocol::MessageId, + req: protocol::ToEnvoyRequestStart, + ) -> Result<()> { + if !self.has_actor(&req.actor_id).await { + self + .send_unknown_actor_response(ws_stream, &message_id.gateway_id, &message_id.request_id) + .await?; + return Ok(()); + } + + let key = tunnel_key(&message_id.gateway_id, &message_id.request_id); + let state = TunnelRequestState { + gateway_id: message_id.gateway_id, + request_id: message_id.request_id, + next_message_index: 0, + kind: TunnelRequestKind::Http(HttpRequestState { + actor_id: req.actor_id, + _method: req.method, + path: req.path, + _headers: req.headers, + body: req.body.unwrap_or_default(), + stream: req.stream, + }), + }; + let should_dispatch = matches!(&state.kind, TunnelRequestKind::Http(http) if !http.stream); + + self.tunnel_requests.lock().await.insert(key.clone(), state); + + if should_dispatch { + self.dispatch_http_request(ws_stream, &key).await?; + } + + Ok(()) + } + + async fn handle_request_chunk( + &self, + ws_stream: &mut WsStream, + message_id: protocol::MessageId, + chunk: protocol::ToEnvoyRequestChunk, + ) -> Result<()> { + let key = tunnel_key(&message_id.gateway_id, &message_id.request_id); + let mut should_dispatch = false; + + { + let mut requests = self.tunnel_requests.lock().await; + if let Some(TunnelRequestState { + kind: TunnelRequestKind::Http(state), + .. + }) = requests.get_mut(&key) + { + state.body.extend(chunk.body); + should_dispatch = chunk.finish; + } + } + + if should_dispatch { + self.dispatch_http_request(ws_stream, &key).await?; + } + + Ok(()) + } + + async fn handle_request_abort(&self, message_id: protocol::MessageId) { + let key = tunnel_key(&message_id.gateway_id, &message_id.request_id); + self.tunnel_requests.lock().await.remove(&key); + } + + async fn dispatch_http_request(&self, ws_stream: &mut WsStream, key: &str) -> Result<()> { + let request = { + let requests = self.tunnel_requests.lock().await; + let Some(TunnelRequestState { + kind: TunnelRequestKind::Http(state), + .. + }) = requests.get(key) + else { + return Ok(()); + }; + state.clone() + }; + + let response = self.handle_http_request(request).await; + self.send_http_response(ws_stream, key, response).await?; + self.tunnel_requests.lock().await.remove(key); + + Ok(()) + } + + async fn handle_http_request(&self, request: HttpRequestState) -> protocol::ToRivetResponseStart { + match request.path.as_str() { + "/ping" => { + let body = serde_json::to_vec(&serde_json::json!({ + "actorId": request.actor_id, + "status": "ok", + "timestamp": chrono::Utc::now().timestamp_millis(), + })) + .expect("serialize ping response"); + + let headers = HashableMap::from_iter([ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body.len().to_string()), + ]); + + protocol::ToRivetResponseStart { + status: 200, + headers, + body: Some(body), + stream: false, + } + } + "/sleep" => { + self.sleep_actor(&request.actor_id).await; + + let body = b"ok".to_vec(); + let headers = HashableMap::from_iter([ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body.len().to_string()), + ]); + + protocol::ToRivetResponseStart { + status: 200, + headers, + body: Some(body), + stream: false, + } + } + _ => { + let body = b"ok".to_vec(); + let headers = HashableMap::from_iter([( + "content-length".to_string(), + body.len().to_string(), + )]); + + protocol::ToRivetResponseStart { + status: 200, + headers, + body: Some(body), + stream: false, + } + } + } + } + + async fn send_http_response( + &self, + ws_stream: &mut WsStream, + key: &str, + response: protocol::ToRivetResponseStart, + ) -> Result<()> { + self + .send_tunnel_message( + ws_stream, + key, + protocol::ToRivetTunnelMessageKind::ToRivetResponseStart(response), + ) + .await + } + + async fn handle_websocket_open( + &self, + ws_stream: &mut WsStream, + message_id: protocol::MessageId, + open: protocol::ToEnvoyWebSocketOpen, + ) -> Result<()> { + if !self.has_actor(&open.actor_id).await { + self + .send_tunnel_message_direct( + ws_stream, + message_id.gateway_id, + message_id.request_id, + 0, + protocol::ToRivetTunnelMessageKind::ToRivetWebSocketClose( + protocol::ToRivetWebSocketClose { + code: Some(1011), + reason: Some("Actor not found".to_string()), + hibernate: false, + }, + ), + ) + .await?; + return Ok(()); + } + + let key = tunnel_key(&message_id.gateway_id, &message_id.request_id); + self.tunnel_requests.lock().await.insert( + key.clone(), + TunnelRequestState { + gateway_id: message_id.gateway_id, + request_id: message_id.request_id, + next_message_index: 0, + kind: TunnelRequestKind::WebSocket, + }, + ); + + self + .send_tunnel_message( + ws_stream, + &key, + protocol::ToRivetTunnelMessageKind::ToRivetWebSocketOpen( + protocol::ToRivetWebSocketOpen { + can_hibernate: true, + }, + ), + ) + .await + } + + async fn handle_websocket_message( + &self, + ws_stream: &mut WsStream, + message_id: protocol::MessageId, + msg: protocol::ToEnvoyWebSocketMessage, + ) -> Result<()> { + let key = tunnel_key(&message_id.gateway_id, &message_id.request_id); + let exists = self.tunnel_requests.lock().await.contains_key(&key); + if !exists { + return Ok(()); + } + + let text = format!("Echo: {}", String::from_utf8_lossy(&msg.data)); + self + .send_tunnel_message( + ws_stream, + &key, + protocol::ToRivetTunnelMessageKind::ToRivetWebSocketMessage( + protocol::ToRivetWebSocketMessage { + data: text.into_bytes(), + binary: false, + }, + ), + ) + .await?; + + self + .send_tunnel_message( + ws_stream, + &key, + protocol::ToRivetTunnelMessageKind::ToRivetWebSocketMessageAck( + protocol::ToRivetWebSocketMessageAck { + index: message_id.message_index, + }, + ), + ) + .await + } + + async fn handle_websocket_close( + &self, + _ws_stream: &mut WsStream, + message_id: protocol::MessageId, + _close: protocol::ToEnvoyWebSocketClose, + ) -> Result<()> { + let key = tunnel_key(&message_id.gateway_id, &message_id.request_id); + self.tunnel_requests.lock().await.remove(&key); + Ok(()) + } + + async fn send_unknown_actor_response( + &self, + ws_stream: &mut WsStream, + gateway_id: &protocol::GatewayId, + request_id: &protocol::RequestId, + ) -> Result<()> { + let body = b"Actor not found".to_vec(); + let headers = HashableMap::from_iter([ + ("x-rivet-error".to_string(), "envoy.actor_not_found".to_string()), + ("content-length".to_string(), body.len().to_string()), + ]); + + self + .send_tunnel_message_direct( + ws_stream, + gateway_id.clone(), + request_id.clone(), + 0, + protocol::ToRivetTunnelMessageKind::ToRivetResponseStart( + protocol::ToRivetResponseStart { + status: 503, + headers, + body: Some(body), + stream: false, + }, + ), + ) + .await + } + + async fn send_tunnel_message( + &self, + ws_stream: &mut WsStream, + key: &str, + message_kind: protocol::ToRivetTunnelMessageKind, + ) -> Result<()> { + let (gateway_id, request_id, message_index) = { + let mut requests = self.tunnel_requests.lock().await; + let request = requests + .get_mut(key) + .context("missing tunnel request state")?; + let message_index = request.next_message_index; + request.next_message_index = request.next_message_index.wrapping_add(1); + (request.gateway_id, request.request_id, message_index) + }; + + self + .send_tunnel_message_direct(ws_stream, gateway_id, request_id, message_index, message_kind) + .await + } + + async fn send_tunnel_message_direct( + &self, + ws_stream: &mut WsStream, + gateway_id: protocol::GatewayId, + request_id: protocol::RequestId, + message_index: u16, + message_kind: protocol::ToRivetTunnelMessageKind, + ) -> Result<()> { + let payload = protocol::ToRivet::ToRivetTunnelMessage(protocol::ToRivetTunnelMessage { + message_id: protocol::MessageId { + gateway_id, + request_id, + message_index, + }, + message_kind, + }); + + let encoded = utils::encode_to_rivet(payload); + ws_stream.send(Message::Binary(encoded.into())).await?; + Ok(()) + } + async fn send_actor_state_update( &self, actor_id: &str, @@ -945,3 +1409,7 @@ impl Drop for Envoy { tracing::debug!("envoy client dropped, shutdown signaled"); } } + +fn tunnel_key(gateway_id: &protocol::GatewayId, request_id: &protocol::RequestId) -> String { + format!("{}:{}", hex::encode(gateway_id), hex::encode(request_id)) +} diff --git a/engine/sdks/rust/test-envoy/src/lib.rs b/engine/sdks/rust/test-envoy/src/lib.rs new file mode 100644 index 0000000000..095a9eed2a --- /dev/null +++ b/engine/sdks/rust/test-envoy/src/lib.rs @@ -0,0 +1,15 @@ +mod actor; +mod behaviors; +mod envoy; +mod server; +mod utils; + +pub use actor::{ActorConfig, ActorEvent, ActorStartResult, ActorStopResult, KvRequest, TestActor}; +pub use behaviors::{ + CountingCrashActor, CrashNTimesThenSucceedActor, CrashOnStartActor, CustomActor, + CustomActorBuilder, DelayedStartActor, EchoActor, NotifyOnStartActor, SleepImmediatelyActor, + StopImmediatelyActor, TimeoutActor, VerifyInputActor, +}; +pub use envoy::{ActorLifecycleEvent, Envoy, EnvoyBuilder, EnvoyConfig, EnvoyConfigBuilder}; +pub use rivet_envoy_protocol as protocol; +pub use server::run_from_env; diff --git a/engine/sdks/rust/test-envoy/src/main.rs b/engine/sdks/rust/test-envoy/src/main.rs new file mode 100644 index 0000000000..3aa1bcaf3e --- /dev/null +++ b/engine/sdks/rust/test-envoy/src/main.rs @@ -0,0 +1,4 @@ +#[tokio::main] +async fn main() -> anyhow::Result<()> { + rivet_test_envoy::run_from_env().await +} diff --git a/engine/sdks/rust/test-envoy/src/server.rs b/engine/sdks/rust/test-envoy/src/server.rs new file mode 100644 index 0000000000..b0e9540555 --- /dev/null +++ b/engine/sdks/rust/test-envoy/src/server.rs @@ -0,0 +1,256 @@ +use anyhow::{Context, Result}; +use async_stream::stream; +use axum::{ + Router, + body::Bytes, + extract::State, + response::{ + IntoResponse, + Json, + Sse, + sse::{Event, KeepAlive}, + }, + routing::{get, post}, +}; +use rivet_envoy_protocol as protocol; +use serde_json::json; +use std::{convert::Infallible, sync::Arc, time::Duration}; +use tokio::{net::TcpListener, sync::Mutex}; +use tracing_subscriber::EnvFilter; + +use crate::{EchoActor, Envoy, EnvoyBuilder, EnvoyConfig}; + +#[derive(Clone)] +struct Settings { + internal_server_port: u16, + namespace: String, + pool_name: String, + envoy_version: u32, + endpoint: String, + token: String, + autostart_server: bool, + autostart_envoy: bool, + autoconfigure_serverless: bool, +} + +impl Settings { + fn from_env() -> Self { + Self { + internal_server_port: std::env::var("INTERNAL_SERVER_PORT") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(5051), + namespace: std::env::var("RIVET_NAMESPACE").unwrap_or_else(|_| "default".to_string()), + pool_name: std::env::var("RIVET_POOL_NAME").unwrap_or_else(|_| "test-envoy".to_string()), + envoy_version: std::env::var("RIVET_ENVOY_VERSION") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(1), + endpoint: std::env::var("RIVET_ENDPOINT") + .unwrap_or_else(|_| "http://127.0.0.1:6420".to_string()), + token: std::env::var("RIVET_TOKEN").unwrap_or_else(|_| "dev".to_string()), + autostart_server: read_bool_env("AUTOSTART_SERVER", true), + autostart_envoy: read_bool_env("AUTOSTART_ENVOY", false), + autoconfigure_serverless: read_bool_env("AUTOCONFIGURE_SERVERLESS", true), + } + } +} + +#[derive(Clone)] +struct AppState { + settings: Settings, + envoy: Arc>>>, +} + +pub async fn run_from_env() -> Result<()> { + init_tracing(); + + let settings = Settings::from_env(); + let state = AppState { + settings: settings.clone(), + envoy: Arc::new(Mutex::new(None)), + }; + + if settings.autostart_envoy { + let envoy = start_envoy(&settings).await?; + *state.envoy.lock().await = Some(envoy); + } else if settings.autoconfigure_serverless { + auto_configure_serverless(&settings).await?; + } + + let server = if settings.autostart_server { + Some(tokio::spawn(run_http_server(state.clone()))) + } else { + None + }; + + install_signal_handlers(); + + if let Some(server) = server { + server.await.context("http server task failed")??; + } else if settings.autostart_envoy { + std::future::pending::<()>().await; + } + + Ok(()) +} + +async fn run_http_server(state: AppState) -> Result<()> { + let app = Router::new() + .route("/health", get(health)) + .route("/shutdown", get(shutdown)) + .route("/api/rivet/start", post(start_serverless)) + .route("/api/rivet/metadata", get(metadata)) + .with_state(state.clone()); + + let addr = format!("0.0.0.0:{}", state.settings.internal_server_port); + let listener = TcpListener::bind(&addr) + .await + .with_context(|| format!("failed to bind {addr}"))?; + + tracing::info!(port = state.settings.internal_server_port, "internal http server listening"); + + axum::serve(listener, app).await.context("http server failed") +} + +async fn health() -> &'static str { + "ok" +} + +async fn shutdown(State(state): State) -> &'static str { + if let Some(envoy) = state.envoy.lock().await.clone() { + let _ = envoy.shutdown().await; + } + "ok" +} + +async fn metadata() -> Json { + Json(json!({ + "runtime": "rivetkit", + "version": "1", + "envoyProtocolVersion": protocol::PROTOCOL_VERSION, + })) +} + +async fn start_serverless( + State(state): State, + body: Bytes, +) -> impl IntoResponse { + tracing::info!("received serverless start request"); + + let envoy = match start_envoy(&state.settings).await { + Ok(envoy) => envoy, + Err(err) => { + tracing::error!(?err, "failed to start serverless envoy"); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + }; + + if let Err(err) = envoy.start_serverless_actor(body.as_ref()).await { + tracing::error!(?err, "failed to inject serverless start payload"); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + + *state.envoy.lock().await = Some(envoy.clone()); + + let stream = stream! { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + loop { + interval.tick().await; + yield Ok::(Event::default().event("ping").data("")); + } + }; + + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() +} + +async fn start_envoy(settings: &Settings) -> Result> { + let config = EnvoyConfig::builder() + .endpoint(&settings.endpoint) + .token(&settings.token) + .namespace(&settings.namespace) + .pool_name(&settings.pool_name) + .version(settings.envoy_version) + .build()?; + + let envoy = EnvoyBuilder::new(config) + .with_default_actor_behavior(|_config| Box::new(EchoActor::new())) + .build()?; + let envoy = Arc::new(envoy); + + envoy.start().await?; + envoy.wait_ready().await; + + Ok(envoy) +} + +async fn auto_configure_serverless(settings: &Settings) -> Result<()> { + tracing::info!("configuring serverless"); + + let client = reqwest::Client::new(); + let url = format!( + "{}/runner-configs/{}?namespace={}", + settings.endpoint.trim_end_matches('/'), + settings.pool_name, + settings.namespace, + ); + let body = json!({ + "datacenters": { + "default": { + "serverless": { + "url": format!("http://localhost:{}/api/rivet", settings.internal_server_port), + "request_lifespan": 300, + "max_concurrent_actors": 10000, + "max_runners": 10000, + "slots_per_runner": 1 + } + } + } + }); + + let response = client + .put(url) + .bearer_auth(&settings.token) + .json(&body) + .send() + .await + .context("failed to upsert serverless config")?; + + if !response.status().is_success() { + let status = response.status(); + let text = response.text().await.unwrap_or_default(); + anyhow::bail!("serverless config request failed: {}: {}", status, text); + } + + Ok(()) +} + +fn init_tracing() { + let filter = EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("info,rivet_test_envoy=debug,rivet_envoy_client=debug")); + + let _ = tracing_subscriber::fmt() + .with_env_filter(filter) + .with_target(false) + .with_ansi(true) + .try_init(); +} + +fn install_signal_handlers() { + tokio::spawn(async { + if tokio::signal::ctrl_c().await.is_ok() { + tracing::debug!("received stop signal, force exiting in 3s"); + tokio::time::sleep(Duration::from_secs(3)).await; + std::process::exit(0); + } + }); +} + +fn read_bool_env(name: &str, default: bool) -> bool { + match std::env::var(name) { + Ok(value) => value == "1", + Err(_) => default, + } +} diff --git a/engine/sdks/rust/envoy-client/src/utils.rs b/engine/sdks/rust/test-envoy/src/utils.rs similarity index 100% rename from engine/sdks/rust/envoy-client/src/utils.rs rename to engine/sdks/rust/test-envoy/src/utils.rs diff --git a/engine/sdks/schemas/README.md b/engine/sdks/schemas/README.md index 98951367a0..be9e8c9f92 100644 --- a/engine/sdks/schemas/README.md +++ b/engine/sdks/schemas/README.md @@ -2,5 +2,4 @@ These are version BARE files that are used for codegen. -`sdks/rust/runner-protocol/build.rs` generates `sdks/schemas/runner-protocol/`. - +`packages/runner-protocol/build.rs` generates `sdks/schemas/runner-protocol/`. diff --git a/engine/sdks/typescript/test-envoy/package.json b/engine/sdks/typescript/test-envoy/package.json deleted file mode 100644 index ad55c8228f..0000000000 --- a/engine/sdks/typescript/test-envoy/package.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "name": "@rivetkit/engine-test-envoy", - "version": "2.1.11-rc.1", - "type": "module", - "scripts": { - "start": "tsx src/index.ts", - "start:bun": "bun src/index.ts", - "build": "tsup src/index.ts", - "check-types": "tsc --noEmit" - }, - "dependencies": { - "@hono/node-server": "^1.19.1", - "@rivetkit/engine-envoy-client": "workspace:*", - "@rivetkit/engine-envoy-protocol": "workspace:*", - "hono": "^4.7.0", - "pino": "^9.9.5", - "ws": "^8.18.3" - }, - "devDependencies": { - "@types/bun": "^1.3.11", - "@types/node": "^22.18.1", - "@types/ws": "^8.18.1", - "tinybench": "^5.0.1", - "tsup": "^8.5.0", - "tsx": "^4.20.5", - "typescript": "^5.9.2", - "vitest": "^1.6.1" - } -} \ No newline at end of file diff --git a/engine/sdks/typescript/test-envoy/src/index.ts b/engine/sdks/typescript/test-envoy/src/index.ts deleted file mode 100644 index e98c3714a4..0000000000 --- a/engine/sdks/typescript/test-envoy/src/index.ts +++ /dev/null @@ -1,278 +0,0 @@ -import * as protocol from "@rivetkit/engine-envoy-protocol"; -import { ShutdownReason, EnvoyHandle, startEnvoy, startEnvoySync } from "@rivetkit/engine-envoy-client"; -import { Hono, type Context as HonoContext, type Next } from "hono"; -import { streamSSE } from "hono/streaming"; -import type { Logger } from "pino"; -import type WebSocket from "ws"; -import { getLogger } from "./log"; - -const INTERNAL_SERVER_PORT = process.env.INTERNAL_SERVER_PORT - ? Number(process.env.INTERNAL_SERVER_PORT) - : 5051; -const RIVET_NAMESPACE = process.env.RIVET_NAMESPACE ?? "default"; -const RIVET_POOL_NAME = process.env.RIVET_POOL_NAME ?? "test-envoy"; -const RIVET_ENVOY_VERSION = process.env.RIVET_ENVOY_VERSION - ? Number(process.env.RIVET_ENVOY_VERSION) - : 1; -const RIVET_ENDPOINT = process.env.RIVET_ENDPOINT ?? "http://127.0.0.1:6420"; -const RIVET_TOKEN = process.env.RIVET_TOKEN ?? "dev"; -const AUTOSTART_SERVER = (process.env.AUTOSTART_SERVER ?? "1") == "1"; -const AUTOSTART_ENVOY = (process.env.AUTOSTART_ENVOY ?? "0") == "1"; -const AUTOCONFIGURE_SERVERLESS = (process.env.AUTOCONFIGURE_SERVERLESS ?? "1") == "1"; - -let envoy: EnvoyHandle | null = null; -const websocketLastMsgIndexes: Map = new Map(); -const config = { - logger: getLogger(), - version: RIVET_ENVOY_VERSION, - endpoint: RIVET_ENDPOINT, - token: RIVET_TOKEN, - namespace: RIVET_NAMESPACE, - poolName: RIVET_POOL_NAME, - prepopulateActorNames: {}, - fetch: async ( - envoy: EnvoyHandle, - actorId: string, - _gatewayId: ArrayBuffer, - _requestId: ArrayBuffer, - request: Request, - ) => { - getLogger().info( - `Fetch called for actor ${actorId}, URL: ${request.url}`, - ); - const url = new URL(request.url); - if (url.pathname === "/ping") { - // Return the actor ID in response - const responseData = { - actorId, - status: "ok", - timestamp: Date.now(), - }; - - return new Response(JSON.stringify(responseData), { - status: 200, - headers: { "Content-Type": "application/json" }, - }); - } else if (url.pathname === "/sleep") { - envoy.sleepActor(actorId); - - return new Response("ok", { - status: 200, - headers: { "Content-Type": "application/json" }, - }); - } - - return new Response("ok", { status: 200 }); - }, - onActorStart: async ( - _envoy: EnvoyHandle, - _actorId: string, - _generation: number, - _config: protocol.ActorConfig, - ) => { - getLogger().info( - `Actor ${_actorId} started (generation ${_generation})`, - ); - }, - onActorStop: async ( - _envoy: EnvoyHandle, - _actorId: string, - _generation: number, - reason: protocol.StopActorReason, - ) => { - getLogger().info( - `Actor ${_actorId} stopped (generation ${_generation})`, - ); - }, - onShutdown() { }, - websocket: async ( - envoy: EnvoyHandle, - actorId: string, - ws: WebSocket, - _gatewayId: ArrayBuffer, - _requestId: ArrayBuffer, - _request: Request, - ) => { - getLogger().info(`WebSocket connected for actor ${actorId}`); - - // Echo server - send back any messages received - ws.addEventListener("message", (event) => { - const data = event.data; - getLogger().info({ - msg: `WebSocket message from actor ${actorId}`, - data, - index: (event as any).rivetMessageIndex, - }); - - ws.send(`Echo: ${data}`); - - // Ack - const websocketId = Buffer.from( - (event as any).rivetRequestId, - ).toString("base64"); - websocketLastMsgIndexes.set( - websocketId, - (event as any).rivetMessageIndex, - ); - envoy.sendHibernatableWebSocketMessageAck( - (event as any).rivetGatewayId, - (event as any).rivetRequestId, - (event as any).rivetMessageIndex, - ); - }); - - ws.addEventListener("close", () => { - getLogger().info(`WebSocket closed for actor ${actorId}`); - }); - - ws.addEventListener("error", (error) => { - getLogger().error({ - msg: `WebSocket error for actor ${actorId}:`, - error, - }); - }); - }, - hibernatableWebSocket: { - canHibernate() { - return true; - }, - }, -}; - -// Create internal server -const app = new Hono(); - -function loggerMiddleware(logger: Logger) { - return async (c: HonoContext, next: Next) => { - const method = c.req.method; - const path = c.req.path; - const startTime = Date.now(); - - await next(); - - const duration = Date.now() - startTime; - logger.debug({ - msg: "http request", - method, - path, - status: c.res.status, - dt: `${duration}ms`, - reqSize: c.req.header("content-length"), - resSize: c.res.headers.get("content-length"), - userAgent: c.req.header("user-agent"), - }); - }; -} -app.use("*", loggerMiddleware(getLogger())); - -app.get("/health", (c) => { - return c.text("ok"); -}); - -app.get("/shutdown", async (c) => { - envoy?.shutdown(false); - return c.text("ok"); -}); - -app.post("/api/rivet/start", async (c) => { - getLogger().info({ - msg: `Received SSE request`, - }); - - let payload = await c.req.arrayBuffer(); - - return streamSSE(c, async (stream) => { - c.req.raw.signal.addEventListener("abort", () => { - getLogger().debug("SSE aborted"); - }); - - const envoy = startEnvoySync({ - ...config, - }); - - envoy.startServerlessActor(payload); - - while (true) { - if (stream.closed || stream.aborted) break; - - await stream.writeSSE({ event: "ping", data: "" }); - await stream.sleep(1000); - } - }); -}); - -app.get("/api/rivet/metadata", async (c) => { - return c.json({ - // Not actually rivetkit - runtime: "rivetkit", - version: "1", - envoyProtocolVersion: protocol.VERSION, - }); -}); - -if (AUTOSTART_SERVER) { - if (process.versions.bun) { - Bun.serve({ fetch: app.fetch, port: INTERNAL_SERVER_PORT, idleTimeout: 0 }); - } else { - const { serve } = await import("@hono/node-server"); - serve({ fetch: app.fetch, port: INTERNAL_SERVER_PORT }); - } - getLogger().info( - `Internal HTTP server listening on port ${INTERNAL_SERVER_PORT}`, - ); -} - -if (AUTOSTART_ENVOY) { - envoy = await startEnvoy(config); -} else if (AUTOCONFIGURE_SERVERLESS) { - await autoConfigureServerless(); -} - -process.on("SIGTERM", async () => { - getLogger().debug("received SIGTERM, force exiting in 3s"); - - await new Promise(res => setTimeout(res, 3000)); - - process.exit(0); -}); -process.on("SIGINT", async () => { - getLogger().debug("received SIGTERM, force exiting in 3s"); - - await new Promise(res => setTimeout(res, 3000)); - - process.exit(0); -}); - -async function autoConfigureServerless() { - getLogger().info("Configuring serverless"); - - const res = await fetch( - `${RIVET_ENDPOINT}/runner-configs/${RIVET_POOL_NAME}?namespace=${RIVET_NAMESPACE}`, - { - method: "PUT", - headers: { - Authorization: `Bearer ${RIVET_TOKEN}`, - "Content-Type": "application/json", - }, - body: JSON.stringify({ - datacenters: { - default: { - serverless: { - url: `http://localhost:${INTERNAL_SERVER_PORT}/api/rivet`, - request_lifespan: 300, - max_concurrent_actors: 10000, - max_runners: 10000, - slots_per_runner: 1, - }, - }, - }, - }), - }, - ); - - if (!res.ok) { - throw new Error( - `request failed: ${res.statusText} (${res.status}):\n${await res.text()}`, - ); - } -} diff --git a/engine/sdks/typescript/test-envoy/src/log.ts b/engine/sdks/typescript/test-envoy/src/log.ts deleted file mode 100644 index a05ece3819..0000000000 --- a/engine/sdks/typescript/test-envoy/src/log.ts +++ /dev/null @@ -1,175 +0,0 @@ -import { - type LevelWithSilent, - type Logger, - pino, - stdTimeFunctions, -} from "pino"; -import { castToLogValue, formatTimestamp, stringify } from "./logfmt"; - -export type { Logger } from "pino"; - -let baseLogger: Logger | undefined; - -/** Cache of child loggers by logger name. */ -const loggerCache = new Map(); - -export function getPinoLevel(): LevelWithSilent { - // Priority: env > default - return (process.env["LOG_LEVEL"] || "warn") - .toString() - .toLowerCase() as LevelWithSilent; -} - -export function getIncludeTarget(): boolean { - return process.env["LOG_TARGET"] === "1"; -} - -/** - * Configure a custom base logger. - */ -export function configureBaseLogger(logger: Logger): void { - baseLogger = logger; - loggerCache.clear(); -} - -// TODO: This can be simplified in logfmt.ts -function customWrite(level: string, o: any) { - const entries: any = {}; - - // Add timestamp if enabled - if (process.env["LOG_TIMESTAMP"] !== "0" && o.time) { - const date = typeof o.time === "number" ? new Date(o.time) : new Date(); - entries.ts = formatTimestamp(date); - } - - // Add level - entries.level = level.toUpperCase(); - - // Add target if present - if (o.target) { - entries.target = o.target; - } - - // Add message - if (o.msg) { - entries.msg = o.msg; - } - - // Add other properties - for (const [key, value] of Object.entries(o)) { - if ( - key !== "time" && - key !== "level" && - key !== "target" && - key !== "msg" && - key !== "pid" && - key !== "hostname" - ) { - entries[key] = castToLogValue(value); - } - } - - const output = stringify(entries); - console.log(output); -} - -/** - * Configure the default logger with optional log level. - */ -export async function configureDefaultLogger(): Promise { - baseLogger = pino({ - level: getPinoLevel(), - messageKey: "msg", - // Do not include pid/hostname in output - base: {}, - // Keep a string level in the output - formatters: { - level(_label: string, number: number) { - return { level: number }; - }, - }, - timestamp: - process.env["LOG_TIMESTAMP"] !== "0" - ? stdTimeFunctions.epochTime - : false, - browser: { - write: { - fatal: customWrite.bind(null, "fatal"), - error: customWrite.bind(null, "error"), - warn: customWrite.bind(null, "warn"), - info: customWrite.bind(null, "info"), - debug: customWrite.bind(null, "debug"), - trace: customWrite.bind(null, "trace"), - }, - }, - hooks: { - logMethod(inputArgs, _method, level) { - // TODO: This is a hack to not implement our own transport target. We can get better perf if we have our own transport target. - - const levelMap: Record = { - 10: "trace", - 20: "debug", - 30: "info", - 40: "warn", - 50: "error", - 60: "fatal", - }; - const levelName = levelMap[level] || "info"; - const time = - process.env["LOG_TIMESTAMP"] !== "0" - ? Date.now() - : undefined; - // TODO: This can be simplified in logfmt.ts - if (inputArgs.length >= 2) { - const [objOrMsg, msg] = inputArgs; - if (typeof objOrMsg === "object" && objOrMsg !== null) { - customWrite(levelName, { ...objOrMsg, msg, time }); - } else { - customWrite(levelName, { msg: String(objOrMsg), time }); - } - } else if (inputArgs.length === 1) { - const [objOrMsg] = inputArgs; - if (typeof objOrMsg === "object" && objOrMsg !== null) { - customWrite(levelName, { ...objOrMsg, time }); - } else { - customWrite(levelName, { msg: String(objOrMsg), time }); - } - } - }, - }, - }); - - loggerCache.clear(); -} - -/** - * Get or initialize the base logger. - */ -export function getBaseLogger(): Logger { - if (!baseLogger) { - configureDefaultLogger(); - } - return baseLogger!; -} - -/** - * Returns a child logger with `target` bound for the given name. - */ -export function getLogger(name = "default"): Logger { - // Check cache first - const cached = loggerCache.get(name); - if (cached) { - return cached; - } - - // Create - const base = getBaseLogger(); - - // Add target to log if enabled - const child = getIncludeTarget() ? base.child({ target: name }) : base; - - // Cache the logger - loggerCache.set(name, child); - - return child; -} diff --git a/engine/sdks/typescript/test-envoy/src/logfmt.ts b/engine/sdks/typescript/test-envoy/src/logfmt.ts deleted file mode 100644 index 53666bba03..0000000000 --- a/engine/sdks/typescript/test-envoy/src/logfmt.ts +++ /dev/null @@ -1,110 +0,0 @@ -import { inspect } from "node:util"; - -type LogLevel = "TRACE" | "DEBUG" | "INFO" | "WARN" | "ERROR" | "FATAL"; - -interface LoggerConfig { - enableInspect: boolean; -} - -export const LOGGER_CONFIG: LoggerConfig = { - enableInspect: true, -}; - -const LOG_LEVEL_COLORS: Record = { - FATAL: "\x1b[31m", - ERROR: "\x1b[31m", - WARN: "\x1b[33m", - INFO: "\x1b[32m", - DEBUG: "\x1b[36m", - TRACE: "\x1b[36m", -}; - -const RESET_COLOR = "\x1b[0m"; - -export function stringify(data: Record): string { - let line = ""; - const entries = Object.entries(data); - - for (let i = 0; i < entries.length; i++) { - const [key, valueRaw] = entries[i]; - - let isNull = false; - let valueString: string; - if (valueRaw == null) { - isNull = true; - valueString = ""; - } else { - valueString = String(valueRaw); - } - - // Clip value unless specifically the error message - if (valueString.length > 512 && key !== "msg" && key !== "error") { - valueString = `${valueString.slice(0, 512)}...`; - } - - const needsQuoting = - valueString.indexOf(" ") > -1 || valueString.indexOf("=") > -1; - const needsEscaping = - valueString.indexOf('"') > -1 || valueString.indexOf("\\") > -1; - - valueString = valueString.replace(/\n/g, "\\n"); - if (needsEscaping) valueString = valueString.replace(/["\\]/g, "\\$&"); - if (needsQuoting || needsEscaping) valueString = `"${valueString}"`; - if (valueString === "" && !isNull) valueString = '""'; - - // Special message colors - let color = "\x1b[2m"; - if (key === "level") { - const levelColor = LOG_LEVEL_COLORS[valueString as LogLevel]; - if (levelColor) { - color = levelColor; - } - } else if (key === "msg") { - color = "\x1b[32m"; - } - - line += `\x1b[0m\x1b[1m${key}\x1b[0m\x1b[2m=\x1b[0m${color}${valueString}${RESET_COLOR}`; - - if (i !== entries.length - 1) { - line += " "; - } - } - - return line; -} - -export function formatTimestamp(date: Date): string { - const year = date.getUTCFullYear(); - const month = String(date.getUTCMonth() + 1).padStart(2, "0"); - const day = String(date.getUTCDate()).padStart(2, "0"); - const hours = String(date.getUTCHours()).padStart(2, "0"); - const minutes = String(date.getUTCMinutes()).padStart(2, "0"); - const seconds = String(date.getUTCSeconds()).padStart(2, "0"); - const milliseconds = String(date.getUTCMilliseconds()).padStart(3, "0"); - - return `${year}-${month}-${day}T${hours}:${minutes}:${seconds}.${milliseconds}Z`; -} - -export function castToLogValue(v: unknown): unknown { - if ( - typeof v === "string" || - typeof v === "number" || - typeof v === "bigint" || - typeof v === "boolean" || - v === null || - v === undefined - ) { - return v; - } - if (LOGGER_CONFIG.enableInspect) { - return inspect(v, { compact: true, breakLength: Infinity, colors: false }); - } - if (v instanceof Error) { - return String(v); - } - try { - return JSON.stringify(v); - } catch { - return "[cannot stringify]"; - } -} diff --git a/engine/sdks/typescript/test-envoy/tsup.config.ts b/engine/sdks/typescript/test-envoy/tsup.config.ts deleted file mode 100644 index 83c2eb83fc..0000000000 --- a/engine/sdks/typescript/test-envoy/tsup.config.ts +++ /dev/null @@ -1,7 +0,0 @@ -import { defineConfig } from "tsup"; -import defaultConfig from "../../../../tsup.base.ts"; - -export default defineConfig({ - ...defaultConfig, - format: "esm", -}); diff --git a/engine/sdks/typescript/test-runner/Dockerfile b/engine/sdks/typescript/test-runner/Dockerfile deleted file mode 100644 index 1b33c9b319..0000000000 --- a/engine/sdks/typescript/test-runner/Dockerfile +++ /dev/null @@ -1,27 +0,0 @@ -FROM node:22-slim - -WORKDIR /app - -# Install pnpm (match workspace version) and set CI for non-interactive behavior -RUN corepack enable && corepack prepare pnpm@10.13.1 --activate -ENV CI=true - -# Copy and build SDK -COPY package.json pnpm-lock.yaml pnpm-workspace.yaml tsup.base.ts tsconfig.base.json turbo.json . -COPY shared/typescript/virtual-websocket ./shared/typescript/virtual-websocket/ -COPY engine/sdks/typescript/runner-protocol ./engine/sdks/typescript/runner-protocol/ -COPY engine/sdks/typescript/runner ./engine/sdks/typescript/runner/ - -# Install dependencies -COPY engine/sdks/typescript/test-runner/package.json ./engine/sdks/typescript/test-runner/ -RUN pnpm install --no-frozen-lockfile - -# Build the application -COPY engine/sdks/typescript/test-runner/ ./engine/sdks/typescript/test-runner/ -RUN pnpm build -F @rivetkit/engine-test-runner - -# Expose the HTTP server port -EXPOSE 5050 - -# Run the application -CMD ["node", "engine/sdks/typescript/test-runner/dist/index.js"] diff --git a/frontend/cloud.Dockerfile b/frontend/cloud.Dockerfile index 6eaf4bd332..c865133b78 100644 --- a/frontend/cloud.Dockerfile +++ b/frontend/cloud.Dockerfile @@ -15,8 +15,8 @@ COPY frontend/ frontend/ # Copy engine SDK dependencies COPY engine/sdks/typescript/api-full/ engine/sdks/typescript/api-full/ -COPY engine/sdks/typescript/runner/ engine/sdks/typescript/runner/ -COPY engine/sdks/typescript/runner-protocol/ engine/sdks/typescript/runner-protocol/ +COPY rivetkit-typescript/packages/engine-runner/ rivetkit-typescript/packages/engine-runner/ +COPY rivetkit-typescript/packages/engine-runner-protocol/ rivetkit-typescript/packages/engine-runner-protocol/ # Copy rivetkit dependencies COPY rivetkit-typescript/packages/rivetkit/ rivetkit-typescript/packages/rivetkit/ diff --git a/frontend/inspector.Dockerfile b/frontend/inspector.Dockerfile index 670ac72f5c..53d97e1ee7 100644 --- a/frontend/inspector.Dockerfile +++ b/frontend/inspector.Dockerfile @@ -15,8 +15,8 @@ COPY frontend/ frontend/ # Copy engine SDK dependencies COPY engine/sdks/typescript/api-full/ engine/sdks/typescript/api-full/ -COPY engine/sdks/typescript/runner/ engine/sdks/typescript/runner/ -COPY engine/sdks/typescript/runner-protocol/ engine/sdks/typescript/runner-protocol/ +COPY rivetkit-typescript/packages/engine-runner/ rivetkit-typescript/packages/engine-runner/ +COPY rivetkit-typescript/packages/engine-runner-protocol/ rivetkit-typescript/packages/engine-runner-protocol/ # Copy rivetkit dependencies COPY rivetkit-typescript/packages/rivetkit/ rivetkit-typescript/packages/rivetkit/ diff --git a/frontend/ladle.Dockerfile b/frontend/ladle.Dockerfile index 89f85eb486..b3d8959933 100644 --- a/frontend/ladle.Dockerfile +++ b/frontend/ladle.Dockerfile @@ -18,8 +18,8 @@ COPY frontend/public/examples/ frontend/public/examples/ # Copy engine SDK dependencies COPY engine/sdks/typescript/api-full/ engine/sdks/typescript/api-full/ -COPY engine/sdks/typescript/runner/ engine/sdks/typescript/runner/ -COPY engine/sdks/typescript/runner-protocol/ engine/sdks/typescript/runner-protocol/ +COPY rivetkit-typescript/packages/engine-runner/ rivetkit-typescript/packages/engine-runner/ +COPY rivetkit-typescript/packages/engine-runner-protocol/ rivetkit-typescript/packages/engine-runner-protocol/ # Copy rivetkit dependencies COPY rivetkit-typescript/packages/rivetkit/ rivetkit-typescript/packages/rivetkit/ diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 34265d5adf..9a068ca917 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -4,10 +4,9 @@ packages: - engine/sdks/typescript/api-full - engine/sdks/typescript/envoy-client - engine/sdks/typescript/envoy-protocol - - engine/sdks/typescript/runner - engine/sdks/typescript/kv-channel-protocol - - engine/sdks/typescript/runner-protocol - engine/sdks/typescript/test-envoy + - engine/sdks/typescript/test-envoy-native - engine/sdks/typescript/test-runner - engine/tests/* - examples/* diff --git a/rivetkit-swift/Tests/RivetKitClientTests/TestServer.swift b/rivetkit-swift/Tests/RivetKitClientTests/TestServer.swift index d481344a64..69548e3714 100644 --- a/rivetkit-swift/Tests/RivetKitClientTests/TestServer.swift +++ b/rivetkit-swift/Tests/RivetKitClientTests/TestServer.swift @@ -155,11 +155,11 @@ actor TestServer { ["pnpm", "--filter", "@rivetkit/virtual-websocket", "build"] ), ( - repoRoot.appendingPathComponent("engine/sdks/typescript/runner-protocol/dist/mod.js"), + repoRoot.appendingPathComponent("rivetkit-typescript/packages/engine-runner-protocol/dist/mod.js"), ["pnpm", "--filter", "@rivetkit/engine-runner-protocol", "build"] ), ( - repoRoot.appendingPathComponent("engine/sdks/typescript/runner/dist/mod.js"), + repoRoot.appendingPathComponent("rivetkit-typescript/packages/engine-runner/dist/mod.js"), ["pnpm", "--filter", "@rivetkit/engine-runner", "build"] ), ] diff --git a/rivetkit-typescript/packages/engine-runner-protocol/package.json b/rivetkit-typescript/packages/engine-runner-protocol/package.json new file mode 100644 index 0000000000..88f0f43d20 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner-protocol/package.json @@ -0,0 +1,36 @@ +{ + "name": "@rivetkit/engine-runner-protocol", + "version": "2.2.1", + "license": "Apache-2.0", + "type": "module", + "exports": { + ".": { + "import": { + "types": "./dist/index.d.ts", + "default": "./dist/index.js" + }, + "require": { + "types": "./dist/index.d.cts", + "default": "./dist/index.cjs" + } + } + }, + "files": [ + "dist/**/*.js", + "dist/**/*.d.ts" + ], + "scripts": { + "build": "tsup src/index.ts", + "clean": "rm -rf dist", + "check-types": "tsc --noEmit" + }, + "types": "dist/index.d.ts", + "dependencies": { + "@rivetkit/bare-ts": "^0.6.2" + }, + "devDependencies": { + "@types/node": "^20.19.13", + "tsup": "^8.5.0", + "typescript": "^5.9.2" + } +} diff --git a/rivetkit-typescript/packages/engine-runner-protocol/src/index.ts b/rivetkit-typescript/packages/engine-runner-protocol/src/index.ts new file mode 100644 index 0000000000..1ee90bbfd4 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner-protocol/src/index.ts @@ -0,0 +1,2189 @@ +// @generated - post-processed by build.rs + +import * as bare from "@rivetkit/bare-ts" + +const DEFAULT_CONFIG = /* @__PURE__ */ bare.Config({}) + +export type i64 = bigint +export type u16 = number +export type u32 = number +export type u64 = bigint + +export type Id = string + +export function readId(bc: bare.ByteCursor): Id { + return bare.readString(bc) +} + +export function writeId(bc: bare.ByteCursor, x: Id): void { + bare.writeString(bc, x) +} + +export type Json = string + +export function readJson(bc: bare.ByteCursor): Json { + return bare.readString(bc) +} + +export function writeJson(bc: bare.ByteCursor, x: Json): void { + bare.writeString(bc, x) +} + +export type GatewayId = ArrayBuffer + +export function readGatewayId(bc: bare.ByteCursor): GatewayId { + return bare.readFixedData(bc, 4) +} + +export function writeGatewayId(bc: bare.ByteCursor, x: GatewayId): void { + assert(x.byteLength === 4) + bare.writeFixedData(bc, x) +} + +export type RequestId = ArrayBuffer + +export function readRequestId(bc: bare.ByteCursor): RequestId { + return bare.readFixedData(bc, 4) +} + +export function writeRequestId(bc: bare.ByteCursor, x: RequestId): void { + assert(x.byteLength === 4) + bare.writeFixedData(bc, x) +} + +export type MessageIndex = u16 + +export function readMessageIndex(bc: bare.ByteCursor): MessageIndex { + return bare.readU16(bc) +} + +export function writeMessageIndex(bc: bare.ByteCursor, x: MessageIndex): void { + bare.writeU16(bc, x) +} + +/** + * Basic types + */ +export type KvKey = ArrayBuffer + +export function readKvKey(bc: bare.ByteCursor): KvKey { + return bare.readData(bc) +} + +export function writeKvKey(bc: bare.ByteCursor, x: KvKey): void { + bare.writeData(bc, x) +} + +export type KvValue = ArrayBuffer + +export function readKvValue(bc: bare.ByteCursor): KvValue { + return bare.readData(bc) +} + +export function writeKvValue(bc: bare.ByteCursor, x: KvValue): void { + bare.writeData(bc, x) +} + +export type KvMetadata = { + readonly version: ArrayBuffer + readonly updateTs: i64 +} + +export function readKvMetadata(bc: bare.ByteCursor): KvMetadata { + return { + version: bare.readData(bc), + updateTs: bare.readI64(bc), + } +} + +export function writeKvMetadata(bc: bare.ByteCursor, x: KvMetadata): void { + bare.writeData(bc, x.version) + bare.writeI64(bc, x.updateTs) +} + +/** + * Query types + */ +export type KvListAllQuery = null + +export type KvListRangeQuery = { + readonly start: KvKey + readonly end: KvKey + readonly exclusive: boolean +} + +export function readKvListRangeQuery(bc: bare.ByteCursor): KvListRangeQuery { + return { + start: readKvKey(bc), + end: readKvKey(bc), + exclusive: bare.readBool(bc), + } +} + +export function writeKvListRangeQuery(bc: bare.ByteCursor, x: KvListRangeQuery): void { + writeKvKey(bc, x.start) + writeKvKey(bc, x.end) + bare.writeBool(bc, x.exclusive) +} + +export type KvListPrefixQuery = { + readonly key: KvKey +} + +export function readKvListPrefixQuery(bc: bare.ByteCursor): KvListPrefixQuery { + return { + key: readKvKey(bc), + } +} + +export function writeKvListPrefixQuery(bc: bare.ByteCursor, x: KvListPrefixQuery): void { + writeKvKey(bc, x.key) +} + +export type KvListQuery = + | { readonly tag: "KvListAllQuery"; readonly val: KvListAllQuery } + | { readonly tag: "KvListRangeQuery"; readonly val: KvListRangeQuery } + | { readonly tag: "KvListPrefixQuery"; readonly val: KvListPrefixQuery } + +export function readKvListQuery(bc: bare.ByteCursor): KvListQuery { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "KvListAllQuery", val: null } + case 1: + return { tag: "KvListRangeQuery", val: readKvListRangeQuery(bc) } + case 2: + return { tag: "KvListPrefixQuery", val: readKvListPrefixQuery(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeKvListQuery(bc: bare.ByteCursor, x: KvListQuery): void { + switch (x.tag) { + case "KvListAllQuery": { + bare.writeU8(bc, 0) + break + } + case "KvListRangeQuery": { + bare.writeU8(bc, 1) + writeKvListRangeQuery(bc, x.val) + break + } + case "KvListPrefixQuery": { + bare.writeU8(bc, 2) + writeKvListPrefixQuery(bc, x.val) + break + } + } +} + +function read0(bc: bare.ByteCursor): readonly KvKey[] { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readKvKey(bc)] + for (let i = 1; i < len; i++) { + result[i] = readKvKey(bc) + } + return result +} + +function write0(bc: bare.ByteCursor, x: readonly KvKey[]): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeKvKey(bc, x[i]) + } +} + +/** + * Request types + */ +export type KvGetRequest = { + readonly keys: readonly KvKey[] +} + +export function readKvGetRequest(bc: bare.ByteCursor): KvGetRequest { + return { + keys: read0(bc), + } +} + +export function writeKvGetRequest(bc: bare.ByteCursor, x: KvGetRequest): void { + write0(bc, x.keys) +} + +function read1(bc: bare.ByteCursor): boolean | null { + return bare.readBool(bc) ? bare.readBool(bc) : null +} + +function write1(bc: bare.ByteCursor, x: boolean | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + bare.writeBool(bc, x) + } +} + +function read2(bc: bare.ByteCursor): u64 | null { + return bare.readBool(bc) ? bare.readU64(bc) : null +} + +function write2(bc: bare.ByteCursor, x: u64 | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + bare.writeU64(bc, x) + } +} + +export type KvListRequest = { + readonly query: KvListQuery + readonly reverse: boolean | null + readonly limit: u64 | null +} + +export function readKvListRequest(bc: bare.ByteCursor): KvListRequest { + return { + query: readKvListQuery(bc), + reverse: read1(bc), + limit: read2(bc), + } +} + +export function writeKvListRequest(bc: bare.ByteCursor, x: KvListRequest): void { + writeKvListQuery(bc, x.query) + write1(bc, x.reverse) + write2(bc, x.limit) +} + +function read3(bc: bare.ByteCursor): readonly KvValue[] { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readKvValue(bc)] + for (let i = 1; i < len; i++) { + result[i] = readKvValue(bc) + } + return result +} + +function write3(bc: bare.ByteCursor, x: readonly KvValue[]): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeKvValue(bc, x[i]) + } +} + +export type KvPutRequest = { + readonly keys: readonly KvKey[] + readonly values: readonly KvValue[] +} + +export function readKvPutRequest(bc: bare.ByteCursor): KvPutRequest { + return { + keys: read0(bc), + values: read3(bc), + } +} + +export function writeKvPutRequest(bc: bare.ByteCursor, x: KvPutRequest): void { + write0(bc, x.keys) + write3(bc, x.values) +} + +export type KvDeleteRequest = { + readonly keys: readonly KvKey[] +} + +export function readKvDeleteRequest(bc: bare.ByteCursor): KvDeleteRequest { + return { + keys: read0(bc), + } +} + +export function writeKvDeleteRequest(bc: bare.ByteCursor, x: KvDeleteRequest): void { + write0(bc, x.keys) +} + +export type KvDeleteRangeRequest = { + readonly start: KvKey + readonly end: KvKey +} + +export function readKvDeleteRangeRequest(bc: bare.ByteCursor): KvDeleteRangeRequest { + return { + start: readKvKey(bc), + end: readKvKey(bc), + } +} + +export function writeKvDeleteRangeRequest(bc: bare.ByteCursor, x: KvDeleteRangeRequest): void { + writeKvKey(bc, x.start) + writeKvKey(bc, x.end) +} + +export type KvDropRequest = null + +/** + * Response types + */ +export type KvErrorResponse = { + readonly message: string +} + +export function readKvErrorResponse(bc: bare.ByteCursor): KvErrorResponse { + return { + message: bare.readString(bc), + } +} + +export function writeKvErrorResponse(bc: bare.ByteCursor, x: KvErrorResponse): void { + bare.writeString(bc, x.message) +} + +function read4(bc: bare.ByteCursor): readonly KvMetadata[] { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readKvMetadata(bc)] + for (let i = 1; i < len; i++) { + result[i] = readKvMetadata(bc) + } + return result +} + +function write4(bc: bare.ByteCursor, x: readonly KvMetadata[]): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeKvMetadata(bc, x[i]) + } +} + +export type KvGetResponse = { + readonly keys: readonly KvKey[] + readonly values: readonly KvValue[] + readonly metadata: readonly KvMetadata[] +} + +export function readKvGetResponse(bc: bare.ByteCursor): KvGetResponse { + return { + keys: read0(bc), + values: read3(bc), + metadata: read4(bc), + } +} + +export function writeKvGetResponse(bc: bare.ByteCursor, x: KvGetResponse): void { + write0(bc, x.keys) + write3(bc, x.values) + write4(bc, x.metadata) +} + +export type KvListResponse = { + readonly keys: readonly KvKey[] + readonly values: readonly KvValue[] + readonly metadata: readonly KvMetadata[] +} + +export function readKvListResponse(bc: bare.ByteCursor): KvListResponse { + return { + keys: read0(bc), + values: read3(bc), + metadata: read4(bc), + } +} + +export function writeKvListResponse(bc: bare.ByteCursor, x: KvListResponse): void { + write0(bc, x.keys) + write3(bc, x.values) + write4(bc, x.metadata) +} + +export type KvPutResponse = null + +export type KvDeleteResponse = null + +export type KvDropResponse = null + +/** + * Request/Response unions + */ +export type KvRequestData = + | { readonly tag: "KvGetRequest"; readonly val: KvGetRequest } + | { readonly tag: "KvListRequest"; readonly val: KvListRequest } + | { readonly tag: "KvPutRequest"; readonly val: KvPutRequest } + | { readonly tag: "KvDeleteRequest"; readonly val: KvDeleteRequest } + | { readonly tag: "KvDeleteRangeRequest"; readonly val: KvDeleteRangeRequest } + | { readonly tag: "KvDropRequest"; readonly val: KvDropRequest } + +export function readKvRequestData(bc: bare.ByteCursor): KvRequestData { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "KvGetRequest", val: readKvGetRequest(bc) } + case 1: + return { tag: "KvListRequest", val: readKvListRequest(bc) } + case 2: + return { tag: "KvPutRequest", val: readKvPutRequest(bc) } + case 3: + return { tag: "KvDeleteRequest", val: readKvDeleteRequest(bc) } + case 4: + return { tag: "KvDeleteRangeRequest", val: readKvDeleteRangeRequest(bc) } + case 5: + return { tag: "KvDropRequest", val: null } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeKvRequestData(bc: bare.ByteCursor, x: KvRequestData): void { + switch (x.tag) { + case "KvGetRequest": { + bare.writeU8(bc, 0) + writeKvGetRequest(bc, x.val) + break + } + case "KvListRequest": { + bare.writeU8(bc, 1) + writeKvListRequest(bc, x.val) + break + } + case "KvPutRequest": { + bare.writeU8(bc, 2) + writeKvPutRequest(bc, x.val) + break + } + case "KvDeleteRequest": { + bare.writeU8(bc, 3) + writeKvDeleteRequest(bc, x.val) + break + } + case "KvDeleteRangeRequest": { + bare.writeU8(bc, 4) + writeKvDeleteRangeRequest(bc, x.val) + break + } + case "KvDropRequest": { + bare.writeU8(bc, 5) + break + } + } +} + +export type KvResponseData = + | { readonly tag: "KvErrorResponse"; readonly val: KvErrorResponse } + | { readonly tag: "KvGetResponse"; readonly val: KvGetResponse } + | { readonly tag: "KvListResponse"; readonly val: KvListResponse } + | { readonly tag: "KvPutResponse"; readonly val: KvPutResponse } + | { readonly tag: "KvDeleteResponse"; readonly val: KvDeleteResponse } + | { readonly tag: "KvDropResponse"; readonly val: KvDropResponse } + +export function readKvResponseData(bc: bare.ByteCursor): KvResponseData { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "KvErrorResponse", val: readKvErrorResponse(bc) } + case 1: + return { tag: "KvGetResponse", val: readKvGetResponse(bc) } + case 2: + return { tag: "KvListResponse", val: readKvListResponse(bc) } + case 3: + return { tag: "KvPutResponse", val: null } + case 4: + return { tag: "KvDeleteResponse", val: null } + case 5: + return { tag: "KvDropResponse", val: null } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeKvResponseData(bc: bare.ByteCursor, x: KvResponseData): void { + switch (x.tag) { + case "KvErrorResponse": { + bare.writeU8(bc, 0) + writeKvErrorResponse(bc, x.val) + break + } + case "KvGetResponse": { + bare.writeU8(bc, 1) + writeKvGetResponse(bc, x.val) + break + } + case "KvListResponse": { + bare.writeU8(bc, 2) + writeKvListResponse(bc, x.val) + break + } + case "KvPutResponse": { + bare.writeU8(bc, 3) + break + } + case "KvDeleteResponse": { + bare.writeU8(bc, 4) + break + } + case "KvDropResponse": { + bare.writeU8(bc, 5) + break + } + } +} + +/** + * Core + */ +export enum StopCode { + Ok = "Ok", + Error = "Error", +} + +export function readStopCode(bc: bare.ByteCursor): StopCode { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return StopCode.Ok + case 1: + return StopCode.Error + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeStopCode(bc: bare.ByteCursor, x: StopCode): void { + switch (x) { + case StopCode.Ok: { + bare.writeU8(bc, 0) + break + } + case StopCode.Error: { + bare.writeU8(bc, 1) + break + } + } +} + +export type ActorName = { + readonly metadata: Json +} + +export function readActorName(bc: bare.ByteCursor): ActorName { + return { + metadata: readJson(bc), + } +} + +export function writeActorName(bc: bare.ByteCursor, x: ActorName): void { + writeJson(bc, x.metadata) +} + +function read5(bc: bare.ByteCursor): string | null { + return bare.readBool(bc) ? bare.readString(bc) : null +} + +function write5(bc: bare.ByteCursor, x: string | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + bare.writeString(bc, x) + } +} + +function read6(bc: bare.ByteCursor): ArrayBuffer | null { + return bare.readBool(bc) ? bare.readData(bc) : null +} + +function write6(bc: bare.ByteCursor, x: ArrayBuffer | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + bare.writeData(bc, x) + } +} + +export type ActorConfig = { + readonly name: string + readonly key: string | null + readonly createTs: i64 + readonly input: ArrayBuffer | null +} + +export function readActorConfig(bc: bare.ByteCursor): ActorConfig { + return { + name: bare.readString(bc), + key: read5(bc), + createTs: bare.readI64(bc), + input: read6(bc), + } +} + +export function writeActorConfig(bc: bare.ByteCursor, x: ActorConfig): void { + bare.writeString(bc, x.name) + write5(bc, x.key) + bare.writeI64(bc, x.createTs) + write6(bc, x.input) +} + +export type ActorCheckpoint = { + readonly actorId: Id + readonly generation: u32 + readonly index: i64 +} + +export function readActorCheckpoint(bc: bare.ByteCursor): ActorCheckpoint { + return { + actorId: readId(bc), + generation: bare.readU32(bc), + index: bare.readI64(bc), + } +} + +export function writeActorCheckpoint(bc: bare.ByteCursor, x: ActorCheckpoint): void { + writeId(bc, x.actorId) + bare.writeU32(bc, x.generation) + bare.writeI64(bc, x.index) +} + +/** + * Intent + */ +export type ActorIntentSleep = null + +export type ActorIntentStop = null + +export type ActorIntent = + | { readonly tag: "ActorIntentSleep"; readonly val: ActorIntentSleep } + | { readonly tag: "ActorIntentStop"; readonly val: ActorIntentStop } + +export function readActorIntent(bc: bare.ByteCursor): ActorIntent { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ActorIntentSleep", val: null } + case 1: + return { tag: "ActorIntentStop", val: null } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeActorIntent(bc: bare.ByteCursor, x: ActorIntent): void { + switch (x.tag) { + case "ActorIntentSleep": { + bare.writeU8(bc, 0) + break + } + case "ActorIntentStop": { + bare.writeU8(bc, 1) + break + } + } +} + +/** + * State + */ +export type ActorStateRunning = null + +export type ActorStateStopped = { + readonly code: StopCode + readonly message: string | null +} + +export function readActorStateStopped(bc: bare.ByteCursor): ActorStateStopped { + return { + code: readStopCode(bc), + message: read5(bc), + } +} + +export function writeActorStateStopped(bc: bare.ByteCursor, x: ActorStateStopped): void { + writeStopCode(bc, x.code) + write5(bc, x.message) +} + +export type ActorState = + | { readonly tag: "ActorStateRunning"; readonly val: ActorStateRunning } + | { readonly tag: "ActorStateStopped"; readonly val: ActorStateStopped } + +export function readActorState(bc: bare.ByteCursor): ActorState { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ActorStateRunning", val: null } + case 1: + return { tag: "ActorStateStopped", val: readActorStateStopped(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeActorState(bc: bare.ByteCursor, x: ActorState): void { + switch (x.tag) { + case "ActorStateRunning": { + bare.writeU8(bc, 0) + break + } + case "ActorStateStopped": { + bare.writeU8(bc, 1) + writeActorStateStopped(bc, x.val) + break + } + } +} + +/** + * MARK: Events + */ +export type EventActorIntent = { + readonly intent: ActorIntent +} + +export function readEventActorIntent(bc: bare.ByteCursor): EventActorIntent { + return { + intent: readActorIntent(bc), + } +} + +export function writeEventActorIntent(bc: bare.ByteCursor, x: EventActorIntent): void { + writeActorIntent(bc, x.intent) +} + +export type EventActorStateUpdate = { + readonly state: ActorState +} + +export function readEventActorStateUpdate(bc: bare.ByteCursor): EventActorStateUpdate { + return { + state: readActorState(bc), + } +} + +export function writeEventActorStateUpdate(bc: bare.ByteCursor, x: EventActorStateUpdate): void { + writeActorState(bc, x.state) +} + +function read7(bc: bare.ByteCursor): i64 | null { + return bare.readBool(bc) ? bare.readI64(bc) : null +} + +function write7(bc: bare.ByteCursor, x: i64 | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + bare.writeI64(bc, x) + } +} + +export type EventActorSetAlarm = { + readonly alarmTs: i64 | null +} + +export function readEventActorSetAlarm(bc: bare.ByteCursor): EventActorSetAlarm { + return { + alarmTs: read7(bc), + } +} + +export function writeEventActorSetAlarm(bc: bare.ByteCursor, x: EventActorSetAlarm): void { + write7(bc, x.alarmTs) +} + +export type Event = + | { readonly tag: "EventActorIntent"; readonly val: EventActorIntent } + | { readonly tag: "EventActorStateUpdate"; readonly val: EventActorStateUpdate } + | { readonly tag: "EventActorSetAlarm"; readonly val: EventActorSetAlarm } + +export function readEvent(bc: bare.ByteCursor): Event { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "EventActorIntent", val: readEventActorIntent(bc) } + case 1: + return { tag: "EventActorStateUpdate", val: readEventActorStateUpdate(bc) } + case 2: + return { tag: "EventActorSetAlarm", val: readEventActorSetAlarm(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeEvent(bc: bare.ByteCursor, x: Event): void { + switch (x.tag) { + case "EventActorIntent": { + bare.writeU8(bc, 0) + writeEventActorIntent(bc, x.val) + break + } + case "EventActorStateUpdate": { + bare.writeU8(bc, 1) + writeEventActorStateUpdate(bc, x.val) + break + } + case "EventActorSetAlarm": { + bare.writeU8(bc, 2) + writeEventActorSetAlarm(bc, x.val) + break + } + } +} + +export type EventWrapper = { + readonly checkpoint: ActorCheckpoint + readonly inner: Event +} + +export function readEventWrapper(bc: bare.ByteCursor): EventWrapper { + return { + checkpoint: readActorCheckpoint(bc), + inner: readEvent(bc), + } +} + +export function writeEventWrapper(bc: bare.ByteCursor, x: EventWrapper): void { + writeActorCheckpoint(bc, x.checkpoint) + writeEvent(bc, x.inner) +} + +export type HibernatingRequest = { + readonly gatewayId: GatewayId + readonly requestId: RequestId +} + +export function readHibernatingRequest(bc: bare.ByteCursor): HibernatingRequest { + return { + gatewayId: readGatewayId(bc), + requestId: readRequestId(bc), + } +} + +export function writeHibernatingRequest(bc: bare.ByteCursor, x: HibernatingRequest): void { + writeGatewayId(bc, x.gatewayId) + writeRequestId(bc, x.requestId) +} + +function read8(bc: bare.ByteCursor): readonly HibernatingRequest[] { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readHibernatingRequest(bc)] + for (let i = 1; i < len; i++) { + result[i] = readHibernatingRequest(bc) + } + return result +} + +function write8(bc: bare.ByteCursor, x: readonly HibernatingRequest[]): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeHibernatingRequest(bc, x[i]) + } +} + +export type CommandStartActor = { + readonly config: ActorConfig + readonly hibernatingRequests: readonly HibernatingRequest[] +} + +export function readCommandStartActor(bc: bare.ByteCursor): CommandStartActor { + return { + config: readActorConfig(bc), + hibernatingRequests: read8(bc), + } +} + +export function writeCommandStartActor(bc: bare.ByteCursor, x: CommandStartActor): void { + writeActorConfig(bc, x.config) + write8(bc, x.hibernatingRequests) +} + +export type CommandStopActor = null + +export type Command = + | { readonly tag: "CommandStartActor"; readonly val: CommandStartActor } + | { readonly tag: "CommandStopActor"; readonly val: CommandStopActor } + +export function readCommand(bc: bare.ByteCursor): Command { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "CommandStartActor", val: readCommandStartActor(bc) } + case 1: + return { tag: "CommandStopActor", val: null } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeCommand(bc: bare.ByteCursor, x: Command): void { + switch (x.tag) { + case "CommandStartActor": { + bare.writeU8(bc, 0) + writeCommandStartActor(bc, x.val) + break + } + case "CommandStopActor": { + bare.writeU8(bc, 1) + break + } + } +} + +export type CommandWrapper = { + readonly checkpoint: ActorCheckpoint + readonly inner: Command +} + +export function readCommandWrapper(bc: bare.ByteCursor): CommandWrapper { + return { + checkpoint: readActorCheckpoint(bc), + inner: readCommand(bc), + } +} + +export function writeCommandWrapper(bc: bare.ByteCursor, x: CommandWrapper): void { + writeActorCheckpoint(bc, x.checkpoint) + writeCommand(bc, x.inner) +} + +/** + * We redeclare this so its top level + */ +export type ActorCommandKeyData = + | { readonly tag: "CommandStartActor"; readonly val: CommandStartActor } + | { readonly tag: "CommandStopActor"; readonly val: CommandStopActor } + +export function readActorCommandKeyData(bc: bare.ByteCursor): ActorCommandKeyData { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "CommandStartActor", val: readCommandStartActor(bc) } + case 1: + return { tag: "CommandStopActor", val: null } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeActorCommandKeyData(bc: bare.ByteCursor, x: ActorCommandKeyData): void { + switch (x.tag) { + case "CommandStartActor": { + bare.writeU8(bc, 0) + writeCommandStartActor(bc, x.val) + break + } + case "CommandStopActor": { + bare.writeU8(bc, 1) + break + } + } +} + +export function encodeActorCommandKeyData(x: ActorCommandKeyData, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeActorCommandKeyData(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeActorCommandKeyData(bytes: Uint8Array): ActorCommandKeyData { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readActorCommandKeyData(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + +export type MessageId = { + /** + * Globally unique ID + */ + readonly gatewayId: GatewayId + /** + * Unique ID to the gateway + */ + readonly requestId: RequestId + /** + * Unique ID to the request + */ + readonly messageIndex: MessageIndex +} + +export function readMessageId(bc: bare.ByteCursor): MessageId { + return { + gatewayId: readGatewayId(bc), + requestId: readRequestId(bc), + messageIndex: readMessageIndex(bc), + } +} + +export function writeMessageId(bc: bare.ByteCursor, x: MessageId): void { + writeGatewayId(bc, x.gatewayId) + writeRequestId(bc, x.requestId) + writeMessageIndex(bc, x.messageIndex) +} + +function read9(bc: bare.ByteCursor): ReadonlyMap { + const len = bare.readUintSafe(bc) + const result = new Map() + for (let i = 0; i < len; i++) { + const offset = bc.offset + const key = bare.readString(bc) + if (result.has(key)) { + bc.offset = offset + throw new bare.BareError(offset, "duplicated key") + } + result.set(key, bare.readString(bc)) + } + return result +} + +function write9(bc: bare.ByteCursor, x: ReadonlyMap): void { + bare.writeUintSafe(bc, x.size) + for (const kv of x) { + bare.writeString(bc, kv[0]) + bare.writeString(bc, kv[1]) + } +} + +/** + * HTTP + */ +export type ToClientRequestStart = { + readonly actorId: Id + readonly method: string + readonly path: string + readonly headers: ReadonlyMap + readonly body: ArrayBuffer | null + readonly stream: boolean +} + +export function readToClientRequestStart(bc: bare.ByteCursor): ToClientRequestStart { + return { + actorId: readId(bc), + method: bare.readString(bc), + path: bare.readString(bc), + headers: read9(bc), + body: read6(bc), + stream: bare.readBool(bc), + } +} + +export function writeToClientRequestStart(bc: bare.ByteCursor, x: ToClientRequestStart): void { + writeId(bc, x.actorId) + bare.writeString(bc, x.method) + bare.writeString(bc, x.path) + write9(bc, x.headers) + write6(bc, x.body) + bare.writeBool(bc, x.stream) +} + +export type ToClientRequestChunk = { + readonly body: ArrayBuffer + readonly finish: boolean +} + +export function readToClientRequestChunk(bc: bare.ByteCursor): ToClientRequestChunk { + return { + body: bare.readData(bc), + finish: bare.readBool(bc), + } +} + +export function writeToClientRequestChunk(bc: bare.ByteCursor, x: ToClientRequestChunk): void { + bare.writeData(bc, x.body) + bare.writeBool(bc, x.finish) +} + +export type ToClientRequestAbort = null + +export type ToServerResponseStart = { + readonly status: u16 + readonly headers: ReadonlyMap + readonly body: ArrayBuffer | null + readonly stream: boolean +} + +export function readToServerResponseStart(bc: bare.ByteCursor): ToServerResponseStart { + return { + status: bare.readU16(bc), + headers: read9(bc), + body: read6(bc), + stream: bare.readBool(bc), + } +} + +export function writeToServerResponseStart(bc: bare.ByteCursor, x: ToServerResponseStart): void { + bare.writeU16(bc, x.status) + write9(bc, x.headers) + write6(bc, x.body) + bare.writeBool(bc, x.stream) +} + +export type ToServerResponseChunk = { + readonly body: ArrayBuffer + readonly finish: boolean +} + +export function readToServerResponseChunk(bc: bare.ByteCursor): ToServerResponseChunk { + return { + body: bare.readData(bc), + finish: bare.readBool(bc), + } +} + +export function writeToServerResponseChunk(bc: bare.ByteCursor, x: ToServerResponseChunk): void { + bare.writeData(bc, x.body) + bare.writeBool(bc, x.finish) +} + +export type ToServerResponseAbort = null + +/** + * WebSocket + */ +export type ToClientWebSocketOpen = { + readonly actorId: Id + readonly path: string + readonly headers: ReadonlyMap +} + +export function readToClientWebSocketOpen(bc: bare.ByteCursor): ToClientWebSocketOpen { + return { + actorId: readId(bc), + path: bare.readString(bc), + headers: read9(bc), + } +} + +export function writeToClientWebSocketOpen(bc: bare.ByteCursor, x: ToClientWebSocketOpen): void { + writeId(bc, x.actorId) + bare.writeString(bc, x.path) + write9(bc, x.headers) +} + +export type ToClientWebSocketMessage = { + readonly data: ArrayBuffer + readonly binary: boolean +} + +export function readToClientWebSocketMessage(bc: bare.ByteCursor): ToClientWebSocketMessage { + return { + data: bare.readData(bc), + binary: bare.readBool(bc), + } +} + +export function writeToClientWebSocketMessage(bc: bare.ByteCursor, x: ToClientWebSocketMessage): void { + bare.writeData(bc, x.data) + bare.writeBool(bc, x.binary) +} + +function read10(bc: bare.ByteCursor): u16 | null { + return bare.readBool(bc) ? bare.readU16(bc) : null +} + +function write10(bc: bare.ByteCursor, x: u16 | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + bare.writeU16(bc, x) + } +} + +export type ToClientWebSocketClose = { + readonly code: u16 | null + readonly reason: string | null +} + +export function readToClientWebSocketClose(bc: bare.ByteCursor): ToClientWebSocketClose { + return { + code: read10(bc), + reason: read5(bc), + } +} + +export function writeToClientWebSocketClose(bc: bare.ByteCursor, x: ToClientWebSocketClose): void { + write10(bc, x.code) + write5(bc, x.reason) +} + +export type ToServerWebSocketOpen = { + readonly canHibernate: boolean +} + +export function readToServerWebSocketOpen(bc: bare.ByteCursor): ToServerWebSocketOpen { + return { + canHibernate: bare.readBool(bc), + } +} + +export function writeToServerWebSocketOpen(bc: bare.ByteCursor, x: ToServerWebSocketOpen): void { + bare.writeBool(bc, x.canHibernate) +} + +export type ToServerWebSocketMessage = { + readonly data: ArrayBuffer + readonly binary: boolean +} + +export function readToServerWebSocketMessage(bc: bare.ByteCursor): ToServerWebSocketMessage { + return { + data: bare.readData(bc), + binary: bare.readBool(bc), + } +} + +export function writeToServerWebSocketMessage(bc: bare.ByteCursor, x: ToServerWebSocketMessage): void { + bare.writeData(bc, x.data) + bare.writeBool(bc, x.binary) +} + +export type ToServerWebSocketMessageAck = { + readonly index: MessageIndex +} + +export function readToServerWebSocketMessageAck(bc: bare.ByteCursor): ToServerWebSocketMessageAck { + return { + index: readMessageIndex(bc), + } +} + +export function writeToServerWebSocketMessageAck(bc: bare.ByteCursor, x: ToServerWebSocketMessageAck): void { + writeMessageIndex(bc, x.index) +} + +export type ToServerWebSocketClose = { + readonly code: u16 | null + readonly reason: string | null + readonly hibernate: boolean +} + +export function readToServerWebSocketClose(bc: bare.ByteCursor): ToServerWebSocketClose { + return { + code: read10(bc), + reason: read5(bc), + hibernate: bare.readBool(bc), + } +} + +export function writeToServerWebSocketClose(bc: bare.ByteCursor, x: ToServerWebSocketClose): void { + write10(bc, x.code) + write5(bc, x.reason) + bare.writeBool(bc, x.hibernate) +} + +/** + * To Server + */ +export type ToServerTunnelMessageKind = + /** + * HTTP + */ + | { readonly tag: "ToServerResponseStart"; readonly val: ToServerResponseStart } + | { readonly tag: "ToServerResponseChunk"; readonly val: ToServerResponseChunk } + | { readonly tag: "ToServerResponseAbort"; readonly val: ToServerResponseAbort } + /** + * WebSocket + */ + | { readonly tag: "ToServerWebSocketOpen"; readonly val: ToServerWebSocketOpen } + | { readonly tag: "ToServerWebSocketMessage"; readonly val: ToServerWebSocketMessage } + | { readonly tag: "ToServerWebSocketMessageAck"; readonly val: ToServerWebSocketMessageAck } + | { readonly tag: "ToServerWebSocketClose"; readonly val: ToServerWebSocketClose } + +export function readToServerTunnelMessageKind(bc: bare.ByteCursor): ToServerTunnelMessageKind { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToServerResponseStart", val: readToServerResponseStart(bc) } + case 1: + return { tag: "ToServerResponseChunk", val: readToServerResponseChunk(bc) } + case 2: + return { tag: "ToServerResponseAbort", val: null } + case 3: + return { tag: "ToServerWebSocketOpen", val: readToServerWebSocketOpen(bc) } + case 4: + return { tag: "ToServerWebSocketMessage", val: readToServerWebSocketMessage(bc) } + case 5: + return { tag: "ToServerWebSocketMessageAck", val: readToServerWebSocketMessageAck(bc) } + case 6: + return { tag: "ToServerWebSocketClose", val: readToServerWebSocketClose(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToServerTunnelMessageKind(bc: bare.ByteCursor, x: ToServerTunnelMessageKind): void { + switch (x.tag) { + case "ToServerResponseStart": { + bare.writeU8(bc, 0) + writeToServerResponseStart(bc, x.val) + break + } + case "ToServerResponseChunk": { + bare.writeU8(bc, 1) + writeToServerResponseChunk(bc, x.val) + break + } + case "ToServerResponseAbort": { + bare.writeU8(bc, 2) + break + } + case "ToServerWebSocketOpen": { + bare.writeU8(bc, 3) + writeToServerWebSocketOpen(bc, x.val) + break + } + case "ToServerWebSocketMessage": { + bare.writeU8(bc, 4) + writeToServerWebSocketMessage(bc, x.val) + break + } + case "ToServerWebSocketMessageAck": { + bare.writeU8(bc, 5) + writeToServerWebSocketMessageAck(bc, x.val) + break + } + case "ToServerWebSocketClose": { + bare.writeU8(bc, 6) + writeToServerWebSocketClose(bc, x.val) + break + } + } +} + +export type ToServerTunnelMessage = { + readonly messageId: MessageId + readonly messageKind: ToServerTunnelMessageKind +} + +export function readToServerTunnelMessage(bc: bare.ByteCursor): ToServerTunnelMessage { + return { + messageId: readMessageId(bc), + messageKind: readToServerTunnelMessageKind(bc), + } +} + +export function writeToServerTunnelMessage(bc: bare.ByteCursor, x: ToServerTunnelMessage): void { + writeMessageId(bc, x.messageId) + writeToServerTunnelMessageKind(bc, x.messageKind) +} + +/** + * To Client + */ +export type ToClientTunnelMessageKind = + /** + * HTTP + */ + | { readonly tag: "ToClientRequestStart"; readonly val: ToClientRequestStart } + | { readonly tag: "ToClientRequestChunk"; readonly val: ToClientRequestChunk } + | { readonly tag: "ToClientRequestAbort"; readonly val: ToClientRequestAbort } + /** + * WebSocket + */ + | { readonly tag: "ToClientWebSocketOpen"; readonly val: ToClientWebSocketOpen } + | { readonly tag: "ToClientWebSocketMessage"; readonly val: ToClientWebSocketMessage } + | { readonly tag: "ToClientWebSocketClose"; readonly val: ToClientWebSocketClose } + +export function readToClientTunnelMessageKind(bc: bare.ByteCursor): ToClientTunnelMessageKind { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToClientRequestStart", val: readToClientRequestStart(bc) } + case 1: + return { tag: "ToClientRequestChunk", val: readToClientRequestChunk(bc) } + case 2: + return { tag: "ToClientRequestAbort", val: null } + case 3: + return { tag: "ToClientWebSocketOpen", val: readToClientWebSocketOpen(bc) } + case 4: + return { tag: "ToClientWebSocketMessage", val: readToClientWebSocketMessage(bc) } + case 5: + return { tag: "ToClientWebSocketClose", val: readToClientWebSocketClose(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToClientTunnelMessageKind(bc: bare.ByteCursor, x: ToClientTunnelMessageKind): void { + switch (x.tag) { + case "ToClientRequestStart": { + bare.writeU8(bc, 0) + writeToClientRequestStart(bc, x.val) + break + } + case "ToClientRequestChunk": { + bare.writeU8(bc, 1) + writeToClientRequestChunk(bc, x.val) + break + } + case "ToClientRequestAbort": { + bare.writeU8(bc, 2) + break + } + case "ToClientWebSocketOpen": { + bare.writeU8(bc, 3) + writeToClientWebSocketOpen(bc, x.val) + break + } + case "ToClientWebSocketMessage": { + bare.writeU8(bc, 4) + writeToClientWebSocketMessage(bc, x.val) + break + } + case "ToClientWebSocketClose": { + bare.writeU8(bc, 5) + writeToClientWebSocketClose(bc, x.val) + break + } + } +} + +export type ToClientTunnelMessage = { + readonly messageId: MessageId + readonly messageKind: ToClientTunnelMessageKind +} + +export function readToClientTunnelMessage(bc: bare.ByteCursor): ToClientTunnelMessage { + return { + messageId: readMessageId(bc), + messageKind: readToClientTunnelMessageKind(bc), + } +} + +export function writeToClientTunnelMessage(bc: bare.ByteCursor, x: ToClientTunnelMessage): void { + writeMessageId(bc, x.messageId) + writeToClientTunnelMessageKind(bc, x.messageKind) +} + +export type ToClientPing = { + readonly ts: i64 +} + +export function readToClientPing(bc: bare.ByteCursor): ToClientPing { + return { + ts: bare.readI64(bc), + } +} + +export function writeToClientPing(bc: bare.ByteCursor, x: ToClientPing): void { + bare.writeI64(bc, x.ts) +} + +function read11(bc: bare.ByteCursor): ReadonlyMap { + const len = bare.readUintSafe(bc) + const result = new Map() + for (let i = 0; i < len; i++) { + const offset = bc.offset + const key = bare.readString(bc) + if (result.has(key)) { + bc.offset = offset + throw new bare.BareError(offset, "duplicated key") + } + result.set(key, readActorName(bc)) + } + return result +} + +function write11(bc: bare.ByteCursor, x: ReadonlyMap): void { + bare.writeUintSafe(bc, x.size) + for (const kv of x) { + bare.writeString(bc, kv[0]) + writeActorName(bc, kv[1]) + } +} + +function read12(bc: bare.ByteCursor): ReadonlyMap | null { + return bare.readBool(bc) ? read11(bc) : null +} + +function write12(bc: bare.ByteCursor, x: ReadonlyMap | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + write11(bc, x) + } +} + +function read13(bc: bare.ByteCursor): Json | null { + return bare.readBool(bc) ? readJson(bc) : null +} + +function write13(bc: bare.ByteCursor, x: Json | null): void { + bare.writeBool(bc, x != null) + if (x != null) { + writeJson(bc, x) + } +} + +/** + * MARK: To Server + */ +export type ToServerInit = { + readonly name: string + readonly version: u32 + readonly totalSlots: u32 + readonly prepopulateActorNames: ReadonlyMap | null + readonly metadata: Json | null +} + +export function readToServerInit(bc: bare.ByteCursor): ToServerInit { + return { + name: bare.readString(bc), + version: bare.readU32(bc), + totalSlots: bare.readU32(bc), + prepopulateActorNames: read12(bc), + metadata: read13(bc), + } +} + +export function writeToServerInit(bc: bare.ByteCursor, x: ToServerInit): void { + bare.writeString(bc, x.name) + bare.writeU32(bc, x.version) + bare.writeU32(bc, x.totalSlots) + write12(bc, x.prepopulateActorNames) + write13(bc, x.metadata) +} + +export type ToServerEvents = readonly EventWrapper[] + +export function readToServerEvents(bc: bare.ByteCursor): ToServerEvents { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readEventWrapper(bc)] + for (let i = 1; i < len; i++) { + result[i] = readEventWrapper(bc) + } + return result +} + +export function writeToServerEvents(bc: bare.ByteCursor, x: ToServerEvents): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeEventWrapper(bc, x[i]) + } +} + +function read14(bc: bare.ByteCursor): readonly ActorCheckpoint[] { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readActorCheckpoint(bc)] + for (let i = 1; i < len; i++) { + result[i] = readActorCheckpoint(bc) + } + return result +} + +function write14(bc: bare.ByteCursor, x: readonly ActorCheckpoint[]): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeActorCheckpoint(bc, x[i]) + } +} + +export type ToServerAckCommands = { + readonly lastCommandCheckpoints: readonly ActorCheckpoint[] +} + +export function readToServerAckCommands(bc: bare.ByteCursor): ToServerAckCommands { + return { + lastCommandCheckpoints: read14(bc), + } +} + +export function writeToServerAckCommands(bc: bare.ByteCursor, x: ToServerAckCommands): void { + write14(bc, x.lastCommandCheckpoints) +} + +export type ToServerStopping = null + +export type ToServerPong = { + readonly ts: i64 +} + +export function readToServerPong(bc: bare.ByteCursor): ToServerPong { + return { + ts: bare.readI64(bc), + } +} + +export function writeToServerPong(bc: bare.ByteCursor, x: ToServerPong): void { + bare.writeI64(bc, x.ts) +} + +export type ToServerKvRequest = { + readonly actorId: Id + readonly requestId: u32 + readonly data: KvRequestData +} + +export function readToServerKvRequest(bc: bare.ByteCursor): ToServerKvRequest { + return { + actorId: readId(bc), + requestId: bare.readU32(bc), + data: readKvRequestData(bc), + } +} + +export function writeToServerKvRequest(bc: bare.ByteCursor, x: ToServerKvRequest): void { + writeId(bc, x.actorId) + bare.writeU32(bc, x.requestId) + writeKvRequestData(bc, x.data) +} + +export type ToServer = + | { readonly tag: "ToServerInit"; readonly val: ToServerInit } + | { readonly tag: "ToServerEvents"; readonly val: ToServerEvents } + | { readonly tag: "ToServerAckCommands"; readonly val: ToServerAckCommands } + | { readonly tag: "ToServerStopping"; readonly val: ToServerStopping } + | { readonly tag: "ToServerPong"; readonly val: ToServerPong } + | { readonly tag: "ToServerKvRequest"; readonly val: ToServerKvRequest } + | { readonly tag: "ToServerTunnelMessage"; readonly val: ToServerTunnelMessage } + +export function readToServer(bc: bare.ByteCursor): ToServer { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToServerInit", val: readToServerInit(bc) } + case 1: + return { tag: "ToServerEvents", val: readToServerEvents(bc) } + case 2: + return { tag: "ToServerAckCommands", val: readToServerAckCommands(bc) } + case 3: + return { tag: "ToServerStopping", val: null } + case 4: + return { tag: "ToServerPong", val: readToServerPong(bc) } + case 5: + return { tag: "ToServerKvRequest", val: readToServerKvRequest(bc) } + case 6: + return { tag: "ToServerTunnelMessage", val: readToServerTunnelMessage(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToServer(bc: bare.ByteCursor, x: ToServer): void { + switch (x.tag) { + case "ToServerInit": { + bare.writeU8(bc, 0) + writeToServerInit(bc, x.val) + break + } + case "ToServerEvents": { + bare.writeU8(bc, 1) + writeToServerEvents(bc, x.val) + break + } + case "ToServerAckCommands": { + bare.writeU8(bc, 2) + writeToServerAckCommands(bc, x.val) + break + } + case "ToServerStopping": { + bare.writeU8(bc, 3) + break + } + case "ToServerPong": { + bare.writeU8(bc, 4) + writeToServerPong(bc, x.val) + break + } + case "ToServerKvRequest": { + bare.writeU8(bc, 5) + writeToServerKvRequest(bc, x.val) + break + } + case "ToServerTunnelMessage": { + bare.writeU8(bc, 6) + writeToServerTunnelMessage(bc, x.val) + break + } + } +} + +export function encodeToServer(x: ToServer, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeToServer(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeToServer(bytes: Uint8Array): ToServer { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readToServer(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + +/** + * MARK: To Client + */ +export type ProtocolMetadata = { + readonly runnerLostThreshold: i64 + readonly actorStopThreshold: i64 + readonly serverlessDrainGracePeriod: i64 | null +} + +export function readProtocolMetadata(bc: bare.ByteCursor): ProtocolMetadata { + return { + runnerLostThreshold: bare.readI64(bc), + actorStopThreshold: bare.readI64(bc), + serverlessDrainGracePeriod: read7(bc), + } +} + +export function writeProtocolMetadata(bc: bare.ByteCursor, x: ProtocolMetadata): void { + bare.writeI64(bc, x.runnerLostThreshold) + bare.writeI64(bc, x.actorStopThreshold) + write7(bc, x.serverlessDrainGracePeriod) +} + +export type ToClientInit = { + readonly runnerId: Id + readonly metadata: ProtocolMetadata +} + +export function readToClientInit(bc: bare.ByteCursor): ToClientInit { + return { + runnerId: readId(bc), + metadata: readProtocolMetadata(bc), + } +} + +export function writeToClientInit(bc: bare.ByteCursor, x: ToClientInit): void { + writeId(bc, x.runnerId) + writeProtocolMetadata(bc, x.metadata) +} + +export type ToClientCommands = readonly CommandWrapper[] + +export function readToClientCommands(bc: bare.ByteCursor): ToClientCommands { + const len = bare.readUintSafe(bc) + if (len === 0) { + return [] + } + const result = [readCommandWrapper(bc)] + for (let i = 1; i < len; i++) { + result[i] = readCommandWrapper(bc) + } + return result +} + +export function writeToClientCommands(bc: bare.ByteCursor, x: ToClientCommands): void { + bare.writeUintSafe(bc, x.length) + for (let i = 0; i < x.length; i++) { + writeCommandWrapper(bc, x[i]) + } +} + +export type ToClientAckEvents = { + readonly lastEventCheckpoints: readonly ActorCheckpoint[] +} + +export function readToClientAckEvents(bc: bare.ByteCursor): ToClientAckEvents { + return { + lastEventCheckpoints: read14(bc), + } +} + +export function writeToClientAckEvents(bc: bare.ByteCursor, x: ToClientAckEvents): void { + write14(bc, x.lastEventCheckpoints) +} + +export type ToClientKvResponse = { + readonly requestId: u32 + readonly data: KvResponseData +} + +export function readToClientKvResponse(bc: bare.ByteCursor): ToClientKvResponse { + return { + requestId: bare.readU32(bc), + data: readKvResponseData(bc), + } +} + +export function writeToClientKvResponse(bc: bare.ByteCursor, x: ToClientKvResponse): void { + bare.writeU32(bc, x.requestId) + writeKvResponseData(bc, x.data) +} + +export type ToClient = + | { readonly tag: "ToClientInit"; readonly val: ToClientInit } + | { readonly tag: "ToClientCommands"; readonly val: ToClientCommands } + | { readonly tag: "ToClientAckEvents"; readonly val: ToClientAckEvents } + | { readonly tag: "ToClientKvResponse"; readonly val: ToClientKvResponse } + | { readonly tag: "ToClientTunnelMessage"; readonly val: ToClientTunnelMessage } + | { readonly tag: "ToClientPing"; readonly val: ToClientPing } + +export function readToClient(bc: bare.ByteCursor): ToClient { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToClientInit", val: readToClientInit(bc) } + case 1: + return { tag: "ToClientCommands", val: readToClientCommands(bc) } + case 2: + return { tag: "ToClientAckEvents", val: readToClientAckEvents(bc) } + case 3: + return { tag: "ToClientKvResponse", val: readToClientKvResponse(bc) } + case 4: + return { tag: "ToClientTunnelMessage", val: readToClientTunnelMessage(bc) } + case 5: + return { tag: "ToClientPing", val: readToClientPing(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToClient(bc: bare.ByteCursor, x: ToClient): void { + switch (x.tag) { + case "ToClientInit": { + bare.writeU8(bc, 0) + writeToClientInit(bc, x.val) + break + } + case "ToClientCommands": { + bare.writeU8(bc, 1) + writeToClientCommands(bc, x.val) + break + } + case "ToClientAckEvents": { + bare.writeU8(bc, 2) + writeToClientAckEvents(bc, x.val) + break + } + case "ToClientKvResponse": { + bare.writeU8(bc, 3) + writeToClientKvResponse(bc, x.val) + break + } + case "ToClientTunnelMessage": { + bare.writeU8(bc, 4) + writeToClientTunnelMessage(bc, x.val) + break + } + case "ToClientPing": { + bare.writeU8(bc, 5) + writeToClientPing(bc, x.val) + break + } + } +} + +export function encodeToClient(x: ToClient, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeToClient(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeToClient(bytes: Uint8Array): ToClient { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readToClient(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + +/** + * MARK: To Runner + */ +export type ToRunnerPing = { + readonly gatewayId: GatewayId + readonly requestId: RequestId + readonly ts: i64 +} + +export function readToRunnerPing(bc: bare.ByteCursor): ToRunnerPing { + return { + gatewayId: readGatewayId(bc), + requestId: readRequestId(bc), + ts: bare.readI64(bc), + } +} + +export function writeToRunnerPing(bc: bare.ByteCursor, x: ToRunnerPing): void { + writeGatewayId(bc, x.gatewayId) + writeRequestId(bc, x.requestId) + bare.writeI64(bc, x.ts) +} + +export type ToRunnerClose = null + +/** + * We have to re-declare the entire union since BARE will not generate the + * ser/de for ToClient if it's not a top-level type + */ +export type ToRunner = + | { readonly tag: "ToRunnerPing"; readonly val: ToRunnerPing } + | { readonly tag: "ToRunnerClose"; readonly val: ToRunnerClose } + | { readonly tag: "ToClientCommands"; readonly val: ToClientCommands } + | { readonly tag: "ToClientAckEvents"; readonly val: ToClientAckEvents } + | { readonly tag: "ToClientTunnelMessage"; readonly val: ToClientTunnelMessage } + +export function readToRunner(bc: bare.ByteCursor): ToRunner { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToRunnerPing", val: readToRunnerPing(bc) } + case 1: + return { tag: "ToRunnerClose", val: null } + case 2: + return { tag: "ToClientCommands", val: readToClientCommands(bc) } + case 3: + return { tag: "ToClientAckEvents", val: readToClientAckEvents(bc) } + case 4: + return { tag: "ToClientTunnelMessage", val: readToClientTunnelMessage(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToRunner(bc: bare.ByteCursor, x: ToRunner): void { + switch (x.tag) { + case "ToRunnerPing": { + bare.writeU8(bc, 0) + writeToRunnerPing(bc, x.val) + break + } + case "ToRunnerClose": { + bare.writeU8(bc, 1) + break + } + case "ToClientCommands": { + bare.writeU8(bc, 2) + writeToClientCommands(bc, x.val) + break + } + case "ToClientAckEvents": { + bare.writeU8(bc, 3) + writeToClientAckEvents(bc, x.val) + break + } + case "ToClientTunnelMessage": { + bare.writeU8(bc, 4) + writeToClientTunnelMessage(bc, x.val) + break + } + } +} + +export function encodeToRunner(x: ToRunner, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeToRunner(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeToRunner(bytes: Uint8Array): ToRunner { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readToRunner(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + +/** + * MARK: To Gateway + */ +export type ToGatewayPong = { + readonly requestId: RequestId + readonly ts: i64 +} + +export function readToGatewayPong(bc: bare.ByteCursor): ToGatewayPong { + return { + requestId: readRequestId(bc), + ts: bare.readI64(bc), + } +} + +export function writeToGatewayPong(bc: bare.ByteCursor, x: ToGatewayPong): void { + writeRequestId(bc, x.requestId) + bare.writeI64(bc, x.ts) +} + +export type ToGateway = + | { readonly tag: "ToGatewayPong"; readonly val: ToGatewayPong } + | { readonly tag: "ToServerTunnelMessage"; readonly val: ToServerTunnelMessage } + +export function readToGateway(bc: bare.ByteCursor): ToGateway { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToGatewayPong", val: readToGatewayPong(bc) } + case 1: + return { tag: "ToServerTunnelMessage", val: readToServerTunnelMessage(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToGateway(bc: bare.ByteCursor, x: ToGateway): void { + switch (x.tag) { + case "ToGatewayPong": { + bare.writeU8(bc, 0) + writeToGatewayPong(bc, x.val) + break + } + case "ToServerTunnelMessage": { + bare.writeU8(bc, 1) + writeToServerTunnelMessage(bc, x.val) + break + } + } +} + +export function encodeToGateway(x: ToGateway, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeToGateway(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeToGateway(bytes: Uint8Array): ToGateway { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readToGateway(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + +/** + * MARK: Serverless + */ +export type ToServerlessServerInit = { + readonly runnerId: Id + readonly runnerProtocolVersion: u16 +} + +export function readToServerlessServerInit(bc: bare.ByteCursor): ToServerlessServerInit { + return { + runnerId: readId(bc), + runnerProtocolVersion: bare.readU16(bc), + } +} + +export function writeToServerlessServerInit(bc: bare.ByteCursor, x: ToServerlessServerInit): void { + writeId(bc, x.runnerId) + bare.writeU16(bc, x.runnerProtocolVersion) +} + +export type ToServerlessServer = + | { readonly tag: "ToServerlessServerInit"; readonly val: ToServerlessServerInit } + +export function readToServerlessServer(bc: bare.ByteCursor): ToServerlessServer { + const offset = bc.offset + const tag = bare.readU8(bc) + switch (tag) { + case 0: + return { tag: "ToServerlessServerInit", val: readToServerlessServerInit(bc) } + default: { + bc.offset = offset + throw new bare.BareError(offset, "invalid tag") + } + } +} + +export function writeToServerlessServer(bc: bare.ByteCursor, x: ToServerlessServer): void { + switch (x.tag) { + case "ToServerlessServerInit": { + bare.writeU8(bc, 0) + writeToServerlessServerInit(bc, x.val) + break + } + } +} + +export function encodeToServerlessServer(x: ToServerlessServer, config?: Partial): Uint8Array { + const fullConfig = config != null ? bare.Config(config) : DEFAULT_CONFIG + const bc = new bare.ByteCursor( + new Uint8Array(fullConfig.initialBufferLength), + fullConfig, + ) + writeToServerlessServer(bc, x) + return new Uint8Array(bc.view.buffer, bc.view.byteOffset, bc.offset) +} + +export function decodeToServerlessServer(bytes: Uint8Array): ToServerlessServer { + const bc = new bare.ByteCursor(bytes, DEFAULT_CONFIG) + const result = readToServerlessServer(bc) + if (bc.offset < bc.view.byteLength) { + throw new bare.BareError(bc.offset, "remaining bytes") + } + return result +} + + +function assert(condition: boolean, message?: string): asserts condition { + if (!condition) throw new Error(message ?? "Assertion failed") +} diff --git a/rivetkit-typescript/packages/engine-runner-protocol/tsconfig.json b/rivetkit-typescript/packages/engine-runner-protocol/tsconfig.json new file mode 100644 index 0000000000..306ec40b0d --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner-protocol/tsconfig.json @@ -0,0 +1,9 @@ +{ + "extends": "../../../tsconfig.base.json", + "compilerOptions": { + "declaration": true, + "outDir": "./dist" + }, + "exclude": ["dist", "node_modules"], + "include": ["**/*.ts"] +} diff --git a/rivetkit-typescript/packages/engine-runner-protocol/tsup.config.ts b/rivetkit-typescript/packages/engine-runner-protocol/tsup.config.ts new file mode 100644 index 0000000000..e7d8e5f88d --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner-protocol/tsup.config.ts @@ -0,0 +1,4 @@ +import { defineConfig } from "tsup"; +import defaultConfig from "../../../tsup.base"; + +export default defineConfig(defaultConfig); diff --git a/engine/sdks/typescript/test-envoy/turbo.json b/rivetkit-typescript/packages/engine-runner-protocol/turbo.json similarity index 100% rename from engine/sdks/typescript/test-envoy/turbo.json rename to rivetkit-typescript/packages/engine-runner-protocol/turbo.json diff --git a/rivetkit-typescript/packages/engine-runner/benches/actor-lifecycle.bench.ts b/rivetkit-typescript/packages/engine-runner/benches/actor-lifecycle.bench.ts new file mode 100644 index 0000000000..e4853feb0d --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/benches/actor-lifecycle.bench.ts @@ -0,0 +1,190 @@ +// import { Bench } from "tinybench"; +// import { Runner } from "@/mod"; +// import type { ActorConfig } from "@/mod"; +// import { +// createActor, +// destroyActor, +// setupBenchmarkRunner, +// createPromiseResolver, +// RIVET_ENDPOINT, +// } from "./utils.js"; +// import { afterEach } from "node:test"; +// +// async function runActorLifecycleBenchmark() { +// // Shared state for benchmarks +// let runner: Runner | null = null; +// let namespaceName: string; +// let runnerName: string; +// let createdActors: string[] = []; +// let wakeActorId: string | null = null; +// let stopped: { promise: Promise; resolve: () => void }; +// let started: { promise: Promise; resolve: () => void }; +// +// const bench = new Bench({ +// time: 1000, +// iterations: 10, +// warmupTime: 0, +// warmupIterations: 0, +// throws: true, +// setup: async (task) => { +// // Setup benchmark runner +// console.log("Setting up benchmark runner..."); +// stopped = createPromiseResolver(); +// started = createPromiseResolver(); +// +// const setup = await setupBenchmarkRunner( +// "lifecycle", +// 5054, +// async ( +// _actorId: string, +// _generation: number, +// _config: ActorConfig, +// ) => { +// started.resolve(); +// }, +// async (_actorId: string, _generation: number) => { +// stopped.resolve(); +// }, +// ); +// runner = setup.runner; +// namespaceName = setup.namespaceName; +// runnerName = setup.runnerName; +// +// console.log( +// `Benchmark setup complete. Namespace: ${namespaceName}, Runner: ${runnerName}`, +// ); +// }, +// teardown: async () => { +// if (runner) { +// await runner.shutdown(true); +// } +// +// // Clean up created actors from creation benchmark +// console.log( +// `Cleaning up ${createdActors.length} actors in ${namespaceName}...`, +// ); +// const cleanupActor = createdActors; +// createdActors = []; +// wakeActorId = null; +// for (const actorId of cleanupActor) { +// try { +// await destroyActor(namespaceName, actorId); +// } catch (err) { +// console.warn(`Failed to clean up actor ${actorId}:`, err); +// } +// } +// +// console.log("Benchmark teardown complete!"); +// }, +// }); +// +// bench.add("create actor", async () => { +// const actorResponse = await createActor( +// namespaceName, +// runnerName, +// false, +// ); +// const actorId = actorResponse.actor.actor_id; +// createdActors.push(actorId); +// +// // Ping the actor +// const pingResponse = await fetch(`${RIVET_ENDPOINT}/ping`, { +// method: "GET", +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// if (!pingResponse.ok) throw "Request failed"; +// }); +// +// //bench.add( +// // "wake actor from sleep", +// // async () => { +// // if (!wakeActorId) throw "No wake actor ID"; +// // +// // // Ping the actor +// // const pingResponse = await fetch(`${RIVET_ENDPOINT}/ping`, { +// // method: "GET", +// // headers: { +// // "x-rivet-target": "actor", +// // "x-rivet-actor": wakeActorId, +// // }, +// // }); +// // +// // if (!pingResponse.ok) { +// // console.error( +// // `Ping failed: ${pingResponse.status} ${pingResponse.statusText}`, +// // ); +// // const errorText = await pingResponse.text(); +// // console.error(`Error response: ${errorText}`); +// // throw `Request failed: ${pingResponse.status} ${pingResponse.statusText}`; +// // } +// // }, +// // { +// // beforeEach: async () => { +// // // Reset promise resolvers for this iteration +// // started = createPromiseResolver(); +// // stopped = createPromiseResolver(); +// // +// // // Create the actor that will be used for wake benchmarking +// // console.log('Creating actor'); +// // const wakeActorResponse = await createActor( +// // namespaceName, +// // runnerName, +// // false, +// // "wake-bench-actor", +// // ); +// // wakeActorId = wakeActorResponse.actor.actor_id; +// // createdActors.push(wakeActorId!); +// // +// // // Wait for actor to start +// // await started.promise; +// // +// // // Put actor to sleep initially +// // runner!.sleepActor(wakeActorId!); +// // await stopped.promise; +// // }, +// // }, +// // // TODO(RVT-4979): Add back after sleep cycles fixed +// // //{ +// // // beforeAll: async () => { +// // // // Create the actor that will be used for wake benchmarking +// // // console.log("Creating wake actor..."); +// // // const wakeActorResponse = await createActor( +// // // namespaceName, +// // // runnerName, +// // // false, +// // // "wake-bench-actor", +// // // ); +// // // wakeActorId = wakeActorResponse.actor.actor_id; +// // // createdActors.push(wakeActorId!); +// // // +// // // // Wait for actor to start +// // // await started.promise; +// // // }, +// // // beforeEach: async () => { +// // // console.log("Putting actor to sleep..."); +// // // +// // // // Put actor to sleep initially +// // // stopped = createPromiseResolver(); +// // // runner!.sleepActor(wakeActorId!); +// // // await stopped.promise; +// // // }, +// // //}, +// //); +// +// // Run the benchmark +// console.log("Running benchmarks..."); +// await bench.run(); +// +// // Display results +// console.table(bench.table()); +// +// console.log("Benchmark complete!"); +// } +// +// // Run the benchmark if this file is executed directly +// if (import.meta.url === `file://${process.argv[1]}`) { +// runActorLifecycleBenchmark(); +// } diff --git a/rivetkit-typescript/packages/engine-runner/benches/utils.ts b/rivetkit-typescript/packages/engine-runner/benches/utils.ts new file mode 100644 index 0000000000..df85b7128f --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/benches/utils.ts @@ -0,0 +1,143 @@ +// import { Runner } from "@/mod"; +// import type { RunnerConfig, ActorConfig } from "@/mod"; +// +// export const RIVET_ENDPOINT = +// process.env.RIVET_ENDPOINT ?? "http://localhost:6420"; +// +// export async function createActor( +// namespaceName: string, +// runnerNameSelector: string, +// durable: boolean, +// actorName: string = "bench-actor", +// ): Promise { +// const response = await fetch( +// `${RIVET_ENDPOINT}/actors?namespace=${namespaceName}`, +// { +// method: "POST", +// headers: { +// "Content-Type": "application/json", +// }, +// body: JSON.stringify({ +// name: actorName, +// input: btoa("bench-input"), +// runner_name_selector: runnerNameSelector, +// durable, +// }), +// }, +// ); +// +// if (!response.ok) { +// throw new Error( +// `Failed to create actor: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// +// return response.json(); +// } +// +// export async function destroyActor( +// namespaceName: string, +// actorId: string, +// ): Promise { +// const response = await fetch( +// `${RIVET_ENDPOINT}/actors/${actorId}?namespace=${namespaceName}`, +// { +// method: "DELETE", +// }, +// ); +// +// if (!response.ok) { +// throw new Error( +// `Failed to delete actor: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// } +// +// export async function createNamespace( +// name: string, +// displayName: string, +// ): Promise { +// const response = await fetch(`${RIVET_ENDPOINT}/namespaces`, { +// method: "POST", +// headers: { +// "Content-Type": "application/json", +// }, +// body: JSON.stringify({ +// name, +// display_name: displayName, +// }), +// }); +// +// if (!response.ok) { +// console.warn( +// `Failed to create namespace: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// } +// +// export interface BenchmarkRunnerSetup { +// runner: Runner; +// namespaceName: string; +// runnerName: string; +// } +// +// export async function setupBenchmarkRunner( +// namespaceSuffix: string, +// port: number, +// onActorStart?: ( +// actorId: string, +// generation: number, +// config: ActorConfig, +// ) => Promise, +// onActorStop?: (actorId: string, generation: number) => Promise, +// ): Promise { +// const namespaceName = `bench-${crypto.randomUUID().slice(0, 8)}`; +// const runnerName = `bench-runner`; +// +// let runnerStartedResolver: () => void; +// const runnerStarted = new Promise((resolve) => { +// runnerStartedResolver = resolve; +// }); +// +// const config: RunnerConfig = { +// version: 1, +// endpoint: RIVET_ENDPOINT, +// namespace: namespaceName, +// addresses: { main: { host: "127.0.0.1", port } }, +// totalSlots: 100, +// prepopulateActorNames: [], +// runnerName: runnerName, +// runnerKey: "default", +// onConnected: () => { +// runnerStartedResolver(); +// }, +// onDisconnected: () => {}, +// fetch: async (_actorId: string, request: Request) => { +// return new Response("ok", { status: 200 }); +// }, +// onActorStart: onActorStart || (async () => {}), +// onActorStop: onActorStop || (async () => {}), +// }; +// +// await createNamespace(namespaceName, `Bench ${namespaceSuffix} Namespace`); +// const runner = new Runner(config); +// runner.start(); +// await runnerStarted; +// +// return { runner, namespaceName, runnerName }; +// } +// +// export function createPromiseResolver(): { +// promise: Promise; +// resolve: (value: T) => void; +// reject: (error: any) => void; +// } { +// let resolve: (value: T) => void; +// let reject: (error: any) => void; +// const promise = new Promise((res, rej) => { +// resolve = res; +// reject = rej; +// }); +// return { promise, resolve: resolve!, reject: reject! }; +// } +// diff --git a/rivetkit-typescript/packages/engine-runner/package.json b/rivetkit-typescript/packages/engine-runner/package.json new file mode 100644 index 0000000000..5414a17f70 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/package.json @@ -0,0 +1,43 @@ +{ + "name": "@rivetkit/engine-runner", + "version": "2.2.1", + "type": "module", + "files": [ + "dist", + "src", + "package.json" + ], + "exports": { + "import": { + "types": "./dist/mod.d.ts", + "default": "./dist/mod.js" + }, + "require": { + "types": "./dist/mod.d.cts", + "default": "./dist/mod.cjs" + } + }, + "scripts": { + "build": "tsup src/mod.ts", + "check-types": "tsc --noEmit", + "test": "vitest run", + "test:watch": "vitest", + "bench": "tsx benches/actor-lifecycle.bench.ts" + }, + "dependencies": { + "@rivetkit/virtual-websocket": "workspace:*", + "@rivetkit/engine-runner-protocol": "workspace:*", + "uuid": "^12.0.0", + "pino": "^9.9.5", + "ws": "^8.18.3" + }, + "devDependencies": { + "@types/node": "^22.18.1", + "@types/ws": "^8.18.1", + "tinybench": "^5.0.1", + "tsup": "^8.5.0", + "tsx": "^4.20.5", + "typescript": "^5.9.2", + "vitest": "^1.6.1" + } +} diff --git a/rivetkit-typescript/packages/engine-runner/src/actor.ts b/rivetkit-typescript/packages/engine-runner/src/actor.ts new file mode 100644 index 0000000000..6a1f12455a --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/actor.ts @@ -0,0 +1,217 @@ +import type * as protocol from "@rivetkit/engine-runner-protocol"; +import { logger } from "./log"; +import type { PendingRequest } from "./tunnel"; +import { arraysEqual, idToStr, promiseWithResolvers } from "./utils"; +import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter"; + +export interface ActorConfig { + name: string; + key: string | null; + createTs: bigint; + input: Uint8Array | null; +} + +export class RunnerActor { + actorId: string; + generation: number; + config: ActorConfig; + pendingRequests: Array<{ + gatewayId: protocol.GatewayId; + requestId: protocol.RequestId; + request: PendingRequest; + }> = []; + webSockets: Array<{ + gatewayId: protocol.GatewayId; + requestId: protocol.RequestId; + ws: WebSocketTunnelAdapter; + }> = []; + actorStartPromise: ReturnType>; + + lastCommandIdx: bigint = -1n; + nextEventIdx: bigint = 0n; + eventHistory: protocol.EventWrapper[] = []; + + /** + * If restoreHibernatingRequests has been called. This is used to assert + * that the caller is implemented correctly. + **/ + hibernationRestored: boolean = false; + + /** + * Set when the actor has explicitly requested to stop (e.g. c.destroy()). + * Used to send StopCode.Ok (graceful) vs StopCode.Error (ungraceful) so + * the engine crash policy handles sleepable actors correctly. + **/ + stopIntentSent: boolean = false; + + constructor( + actorId: string, + generation: number, + config: ActorConfig, + /** + * List of hibernating requests provided by the gateway on actor start. + * This represents the WebSocket connections that the gateway knows about. + **/ + public hibernatingRequests: readonly protocol.HibernatingRequest[], + ) { + this.actorId = actorId; + this.generation = generation; + this.config = config; + this.actorStartPromise = promiseWithResolvers(); + } + + // Pending request methods + getPendingRequest( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ): PendingRequest | undefined { + return this.pendingRequests.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + )?.request; + } + + createPendingRequest( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + clientMessageIndex: number, + ) { + const exists = + this.getPendingRequest(gatewayId, requestId) !== undefined; + if (exists) { + logger()?.warn({ + msg: "attempting to set pending request twice, replacing existing", + gatewayId: idToStr(gatewayId), + requestId: idToStr(requestId), + }); + // Delete existing pending request before adding the new one + this.deletePendingRequest(gatewayId, requestId); + } + this.pendingRequests.push({ + gatewayId, + requestId, + request: { + resolve: () => {}, + reject: () => {}, + actorId: this.actorId, + gatewayId: gatewayId, + requestId: requestId, + clientMessageIndex, + }, + }); + logger()?.debug({ + msg: "added pending request", + gatewayId: idToStr(gatewayId), + requestId: idToStr(requestId), + length: this.pendingRequests.length, + }); + } + + createPendingRequestWithStreamController( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + clientMessageIndex: number, + streamController: ReadableStreamDefaultController, + ) { + const exists = + this.getPendingRequest(gatewayId, requestId) !== undefined; + if (exists) { + logger()?.warn({ + msg: "attempting to set pending request twice, replacing existing", + gatewayId: idToStr(gatewayId), + requestId: idToStr(requestId), + }); + // Delete existing pending request before adding the new one + this.deletePendingRequest(gatewayId, requestId); + } + this.pendingRequests.push({ + gatewayId, + requestId, + request: { + resolve: () => {}, + reject: () => {}, + actorId: this.actorId, + gatewayId: gatewayId, + requestId: requestId, + clientMessageIndex, + streamController, + }, + }); + logger()?.debug({ + msg: "added pending request with stream controller", + gatewayId: idToStr(gatewayId), + requestId: idToStr(requestId), + length: this.pendingRequests.length, + }); + } + + deletePendingRequest( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ) { + const index = this.pendingRequests.findIndex( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + if (index !== -1) { + this.pendingRequests.splice(index, 1); + logger()?.debug({ + msg: "removed pending request", + gatewayId: idToStr(gatewayId), + requestId: idToStr(requestId), + length: this.pendingRequests.length, + }); + } + } + + // WebSocket methods + getWebSocket( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ): WebSocketTunnelAdapter | undefined { + return this.webSockets.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + )?.ws; + } + + setWebSocket( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ws: WebSocketTunnelAdapter, + ) { + const exists = this.getWebSocket(gatewayId, requestId) !== undefined; + if (exists) { + logger()?.warn({ msg: "attempting to set websocket twice" }); + return; + } + this.webSockets.push({ gatewayId, requestId, ws }); + } + + deleteWebSocket( + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + ) { + const index = this.webSockets.findIndex( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + if (index !== -1) { + this.webSockets.splice(index, 1); + } + } + + handleAckEvents(lastEventIdx: bigint) { + this.eventHistory = this.eventHistory.filter( + (event) => event.checkpoint.index > lastEventIdx, + ); + } + + recordEvent(eventWrapper: protocol.EventWrapper) { + this.eventHistory.push(eventWrapper); + } +} diff --git a/rivetkit-typescript/packages/engine-runner/src/log.ts b/rivetkit-typescript/packages/engine-runner/src/log.ts new file mode 100644 index 0000000000..e1bf2199e3 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/log.ts @@ -0,0 +1,11 @@ +import type { Logger } from "pino"; + +let LOGGER: Logger | undefined; + +export function setLogger(logger: Logger) { + LOGGER = logger; +} + +export function logger(): Logger | undefined { + return LOGGER; +} diff --git a/rivetkit-typescript/packages/engine-runner/src/mod.ts b/rivetkit-typescript/packages/engine-runner/src/mod.ts new file mode 100644 index 0000000000..724f8ecd49 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/mod.ts @@ -0,0 +1,1927 @@ +import * as protocol from "@rivetkit/engine-runner-protocol"; +import type { Logger } from "pino"; +import type WebSocket from "ws"; +import { type ActorConfig, RunnerActor } from "./actor"; +import { logger, setLogger } from "./log.js"; +import { stringifyToClient, stringifyToServer } from "./stringify"; +import { type HibernatingWebSocketMetadata, Tunnel } from "./tunnel"; +import { + calculateBackoff, + parseWebSocketCloseReason, + stringifyError, + unreachable, +} from "./utils"; +import { importWebSocket } from "./websocket.js"; +import { + v4 as uuidv4, +} from "uuid"; + +export type { HibernatingWebSocketMetadata }; +export { RunnerActor, type ActorConfig }; +export { idToStr } from "./utils"; + +const KV_EXPIRE: number = 30_000; +const PROTOCOL_VERSION: number = 7; + +/** Warn once the backlog significantly exceeds the server's ack batch size. */ +const EVENT_BACKLOG_WARN_THRESHOLD = 10_000; +const SIGNAL_HANDLERS: (() => void | Promise)[] = []; + +export class RunnerShutdownError extends Error { + constructor() { + super("Runner shut down"); + } +} + +export interface RunnerConfig { + logger?: Logger; + version: number; + endpoint: string; + token?: string; + pegboardEndpoint?: string; + pegboardRelayEndpoint?: string; + namespace: string; + totalSlots: number; + runnerName: string; + prepopulateActorNames: Record }>; + metadata?: Record; + onConnected: () => void; + onDisconnected: (code: number, reason: string) => void; + onShutdown: () => void; + + /** Called when receiving a network request. */ + fetch: ( + runner: Runner, + actorId: string, + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + request: Request, + ) => Promise; + + /** + * Called when receiving a WebSocket connection. + * + * All event listeners must be added synchronously inside this function or + * else events may be missed. The open event will fire immediately after + * this function finishes. + * + * Any errors thrown here will disconnect the WebSocket immediately. + * + * While `path` and `headers` are partially redundant to the data in the + * `Request`, they may vary slightly from the actual content of `Request`. + * Prefer to persist the `path` and `headers` properties instead of the + * `Request` itself. + * + * ## Hibernating Web Sockets + * + * ### Implementation Requirements + * + * **Requirement 1: Persist HWS Immediately** + * + * This is responsible for persisting hibernatable WebSockets immediately + * (do not wait for open event). It is not time sensitive to flush the + * connection state. If this fails to persist the HWS, the client's + * WebSocket will be disconnected on next wake in the call to + * `Tunnel::restoreHibernatingRequests` since the connection entry will not + * exist. + * + * **Requirement 2: Persist Message Index On `message`** + * + * In the `message` event listener, this handler must persist the message + * index from the event. The request ID is available at + * `event.rivetRequestId` and message index at `event.rivetMessageIndex`. + * + * The message index should not be flushed immediately. Instead, this + * should: + * + * - Debounce calls to persist the message index + * - After each persist, call + * `Runner::sendHibernatableWebSocketMessageAck` to acknowledge the + * message + * + * This mechanism allows us to buffer messages on the gateway so we can + * batch-persist events on our end on a given interval. + * + * If this fails to persist, then the gateway will replay unacked + * messages when the actor starts again. + * + * **Requirement 3: Remove HWS From Storage On `close`** + * + * This handler should add an event listener for `close` to remove the + * connection from storage. + * + * If the connection remove fails to persist, the close event will be + * called again on the next actor start in + * `Tunnel::restoreHibernatingRequests` since there will be no request for + * the given connection. + * + * ### Restoring Connections + * + * The user of this library is responsible for: + * 1. Loading all persisted hibernatable WebSocket metadata for an actor + * 2. Calling `Runner::restoreHibernatingRequests` with this metadata at + * the end of `onActorStart` + * + * `restoreHibernatingRequests` will restore all connections and attach + * the appropriate event listeners. + * + * ### No Open Event On Restoration + * + * When restoring a HWS, the open event will not be called again. It will + * go straight to the message or close event. + */ + websocket: ( + runner: Runner, + actorId: string, + ws: any, + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + request: Request, + path: string, + headers: Record, + isHibernatable: boolean, + isRestoringHibernatable: boolean, + ) => Promise; + + hibernatableWebSocket: { + /** + * Determines if a WebSocket can continue to live while an actor goes to + * sleep. + */ + canHibernate: ( + actorId: string, + gatewayId: ArrayBuffer, + requestId: ArrayBuffer, + request: Request, + ) => boolean; + }; + + /** + * Called when an actor starts. + * + * This callback is responsible for: + * 1. Initializing the actor instance + * 2. Loading all persisted hibernatable WebSocket metadata for this actor + * 3. Calling `Runner::restoreHibernatingRequests` with the loaded metadata + * to restore hibernatable WebSocket connections + * + * The actor should not be marked as "ready" until after + * `restoreHibernatingRequests` completes to ensure all hibernatable + * connections are fully restored before the actor processes new requests. + */ + onActorStart: ( + actorId: string, + generation: number, + config: ActorConfig, + ) => Promise; + + onActorStop: (actorId: string, generation: number) => Promise; + noAutoShutdown?: boolean; + + /** + * Debug option to inject artificial latency (in ms) into WebSocket + * communication. Messages are queued and delivered in order after the + * configured delay. + * + * @experimental For testing only. + */ + debugLatencyMs?: number; +} + +export interface KvListOptions { + reverse?: boolean; + limit?: number; +} + +interface KvRequestEntry { + actorId: string; + data: protocol.KvRequestData; + resolve: (value: any) => void; + reject: (error: unknown) => void; + sent: boolean; + timestamp: number; +} + +export class Runner { + #config: RunnerConfig; + #runnerKey: string = uuidv4(); + + get config(): RunnerConfig { + return this.#config; + } + + #actors: Map = new Map(); + + // WebSocket + #pegboardWebSocket?: WebSocket; + runnerId?: string; + #started: boolean = false; + #shutdown: boolean = false; + #draining: boolean = false; + #reconnectAttempt: number = 0; + #reconnectTimeout?: NodeJS.Timeout; + + // Protocol metadata + #protocolMetadata?: protocol.ProtocolMetadata; + + // Runner lost threshold management + #runnerLostTimeout?: NodeJS.Timeout; + + // Event storage for resending + #eventBacklogWarned: boolean = false; + + // Command acknowledgment + #ackInterval?: NodeJS.Timeout; + + // KV operations + #nextKvRequestId: number = 0; + #kvRequests: Map = new Map(); + #kvCleanupInterval?: NodeJS.Timeout; + + // Tunnel for HTTP/WebSocket forwarding + #tunnel: Tunnel | undefined; + + // Cached child logger with runner-specific attributes + #logCached?: Logger; + + get log(): Logger | undefined { + if (this.#logCached) return this.#logCached; + + const l = logger(); + if (l) { + // If has connected, create child logger with relevant metadata + // + // Otherwise, return default logger + if (this.runnerId) { + this.#logCached = l.child({ + runnerId: this.runnerId, + }); + return this.#logCached; + } else { + return l; + } + } + + return undefined; + } + + constructor(config: RunnerConfig) { + this.#config = config; + if (this.#config.logger) setLogger(this.#config.logger); + + // Start cleaning up old unsent KV requests every 15 seconds + this.#kvCleanupInterval = setInterval(() => { + try { + this.#cleanupOldKvRequests(); + } catch (err) { + this.log?.error({ + msg: "error cleaning up kv requests", + error: stringifyError(err), + }); + } + }, 15000); // Run every 15 seconds + } + + // MARK: Manage actors + sleepActor(actorId: string, generation?: number) { + const actor = this.getActor(actorId, generation); + if (!actor) return; + + // Keep the actor instance in memory during sleep + this.#sendActorIntent(actorId, actor.generation, "sleep"); + + // NOTE: We do NOT remove the actor from this.#actors here + // The server will send a StopActor command if it wants to fully stop + } + + async stopActor(actorId: string, generation?: number) { + const actor = this.getActor(actorId, generation); + if (!actor) return; + + this.#sendActorIntent(actorId, actor.generation, "stop"); + + // NOTE: We do NOT remove the actor from this.#actors here + // The server will send a StopActor command if it wants to fully stop + } + + /** + * Like stopActor but marks the actor for graceful destruction. + * This ensures the engine destroys the actor instead of sleeping it. + * + * NOTE: If a drain (GoingAway) occurs after this is called but before the + * stop completes, the engine's going_away flag overrides graceful_exit and + * the actor will sleep instead of being destroyed. The destroy intent is + * lost in this race. This is acceptable since the actor will be rescheduled + * elsewhere and can be destroyed on the next wake. + */ + destroyActor(actorId: string, generation?: number) { + const actor = this.getActor(actorId, generation); + if (!actor) return; + + actor.stopIntentSent = true; + this.#sendActorIntent(actorId, actor.generation, "stop"); + } + + async forceStopActor(actorId: string, generation?: number) { + this.log?.debug({ + msg: "force stopping actor", + actorId, + }); + + const actor = this.getActor(actorId, generation); + if (!actor) return; + + // If onActorStop times out, Pegboard will handle this timeout with ACTOR_STOP_THRESHOLD_DURATION_MS + // + // If we receive a request while onActorStop is running, a Service + // Unavailable error will be returned to Guard and the request will be + // retried + try { + await this.#config.onActorStop(actorId, actor.generation); + } catch (err) { + console.error(`Error in onActorStop for actor ${actorId}:`, err); + } + + // Close requests after onActorStop so you can send messages over the tunnel + this.#tunnel?.closeActiveRequests(actor); + + this.#sendActorStateUpdate(actorId, actor.generation, "stopped"); + + // Remove actor after stopping in order to ensure that we can still + // call actions on the runner + this.#removeActor(actorId, generation); + } + + #handleLost() { + this.log?.info({ + msg: "stopping all actors due to runner lost threshold", + }); + + // Remove all remaining kv requests + for (const [_, request] of this.#kvRequests.entries()) { + request.reject(new RunnerShutdownError()); + } + + this.#kvRequests.clear(); + + this.#stopAllActors(); + } + + #stopAllActors() { + const actorIds = Array.from(this.#actors.keys()); + for (const actorId of actorIds) { + this.forceStopActor(actorId).catch((err) => { + this.log?.error({ + msg: "error stopping actor", + actorId, + error: stringifyError(err), + }); + }); + } + } + + getActor(actorId: string, generation?: number): RunnerActor | undefined { + const actor = this.#actors.get(actorId); + if (!actor) { + this.log?.warn({ + msg: "actor not found", + actorId, + }); + return undefined; + } + if (generation !== undefined && actor.generation !== generation) { + this.log?.warn({ + msg: "actor generation mismatch", + actorId, + generation, + }); + return undefined; + } + + return actor; + } + + async getAndWaitForActor( + actorId: string, + generation?: number, + ): Promise { + const actor = this.getActor(actorId, generation); + if (!actor) return; + await actor.actorStartPromise.promise; + return actor; + } + + hasActor(actorId: string, generation?: number): boolean { + const actor = this.#actors.get(actorId); + + return ( + !!actor && + (generation === undefined || actor.generation === generation) + ); + } + + get actors() { + return this.#actors; + } + + // IMPORTANT: Make sure to call stopActiveRequests if calling #removeActor + #removeActor( + actorId: string, + generation?: number, + ): RunnerActor | undefined { + const actor = this.#actors.get(actorId); + if (!actor) { + this.log?.error({ + msg: "actor not found for removal", + actorId, + }); + return undefined; + } + if (generation !== undefined && actor.generation !== generation) { + this.log?.error({ + msg: "actor generation mismatch", + actorId, + generation, + }); + return undefined; + } + + this.#actors.delete(actorId); + + this.log?.info({ + msg: "removed actor", + actorId, + actors: this.#actors.size, + }); + + return actor; + } + + // MARK: Start + async start() { + if (this.#started) throw new Error("Cannot call runner.start twice"); + this.#started = true; + + this.log?.info({ msg: "starting runner" }); + + this.#tunnel = new Tunnel(this); + this.#tunnel.start(); + + try { + await this.#openPegboardWebSocket(); + } catch (error) { + this.#started = false; + throw error; + } + + // When changing SIGTERM/shutdown behavior, update + // website/src/content/docs/actors/versions.mdx (SIGTERM Handling section). + if (!this.#config.noAutoShutdown) { + if (!SIGNAL_HANDLERS.length) { + process.on("SIGTERM", async () => { + this.log?.debug("received SIGTERM"); + + for (const handler of SIGNAL_HANDLERS) { + await handler(); + } + + // TODO: Add back + // process.exit(0); + }); + process.on("SIGINT", async () => { + this.log?.debug("received SIGINT"); + + for (const handler of SIGNAL_HANDLERS) { + await handler(); + } + + // TODO: Add back + // process.exit(0); + }); + + this.log?.debug({ + msg: "added SIGTERM listeners", + }); + } + + SIGNAL_HANDLERS.push(async () => { + const weak = new WeakRef(this); + await weak.deref()?.shutdown(false, false); + }); + } + } + + // MARK: Shutdown + async shutdown(immediate: boolean, exit: boolean = false) { + // Prevent concurrent shutdowns + if (this.#shutdown) { + this.log?.debug({ + msg: "shutdown already in progress, ignoring", + }); + return; + } + this.#shutdown = true; + this.#draining = !immediate; + + this.log?.info({ + msg: "starting shutdown", + immediate, + exit, + }); + + // Clear reconnect timeout + if (this.#reconnectTimeout) { + clearTimeout(this.#reconnectTimeout); + this.#reconnectTimeout = undefined; + } + + // Clear runner lost timeout + if (this.#runnerLostTimeout) { + clearTimeout(this.#runnerLostTimeout); + this.#runnerLostTimeout = undefined; + } + + // Clear ack interval + if (this.#ackInterval) { + clearInterval(this.#ackInterval); + this.#ackInterval = undefined; + } + + // Clear KV cleanup interval + if (this.#kvCleanupInterval) { + clearInterval(this.#kvCleanupInterval); + this.#kvCleanupInterval = undefined; + } + + // Reject all KV requests + for (const request of this.#kvRequests.values()) { + request.reject( + new Error("WebSocket connection closed during shutdown"), + ); + } + this.#kvRequests.clear(); + + // Close WebSocket + const pegboardWebSocket = this.getPegboardWebSocketIfReady(); + if (pegboardWebSocket) { + if (immediate) { + // Stop immediately + pegboardWebSocket.close(1000, "pegboard.runner_shutdown"); + } else { + // Wait for actors to shut down before stopping + try { + this.log?.info({ + msg: "sending stopping message", + readyState: pegboardWebSocket.readyState, + }); + + // Start stopping + // + // The runner workflow will send StopActor commands for all + // actors + this.__sendToServer({ + tag: "ToServerStopping", + val: null, + }); + + const closePromise = new Promise((resolve) => { + if (!pegboardWebSocket) + throw new Error("missing pegboardWebSocket"); + + pegboardWebSocket.addEventListener("close", (ev) => { + this.log?.info({ + msg: "connection closed", + code: ev.code, + reason: ev.reason.toString(), + }); + resolve(); + }); + }); + + // Wait for all actors to stop before closing ws + await this.#waitForActorsToStop(pegboardWebSocket); + + this.log?.info({ + msg: "closing WebSocket", + }); + pegboardWebSocket.close(1000, "pegboard.runner_shutdown"); + + await closePromise; + + this.log?.info({ + msg: "websocket shutdown completed", + }); + } catch (error) { + this.log?.error({ + msg: "error during websocket shutdown:", + error, + }); + pegboardWebSocket.close(); + } + } + } else { + // This is often logged when the serverless SSE stream closes after + // the runner has already shut down + this.log?.debug({ + msg: "no runner WebSocket to shutdown or already closed", + readyState: this.#pegboardWebSocket?.readyState, + }); + } + + // Close tunnel + if (this.#tunnel) { + this.#tunnel.shutdown(); + this.#tunnel = undefined; + } + + this.#config.onShutdown(); + + if (exit) process.exit(0); + } + + /** + * Wait for all actors to stop before proceeding with shutdown. + * + * This method polls every 100ms to check if all actors have been stopped. + * + * It will resolve early if: + * - All actors are stopped + * - The WebSocket connection is closed + * - The shutdown timeout is reached (120 seconds) + * + * When changing this timeout, update + * website/src/content/docs/actors/versions.mdx (SIGTERM Handling section). + */ + async #waitForActorsToStop(ws: WebSocket): Promise { + const shutdownTimeout = 120_000; // 120 seconds + const shutdownCheckInterval = 100; // Check every 100ms + const progressLogInterval = 5_000; // Log progress every 5 seconds + const shutdownStartTs = Date.now(); + let lastProgressLogTs = 0; // Ensure first log happens immediately + + return new Promise((resolve) => { + const checkActors = () => { + const now = Date.now(); + const elapsed = now - shutdownStartTs; + const wsIsClosed = ws.readyState === 2 || ws.readyState === 3; + + if (this.#actors.size === 0) { + this.log?.info({ + msg: "all actors stopped", + elapsed, + }); + return true; + } else if (wsIsClosed) { + this.log?.warn({ + msg: "websocket closed before all actors stopped", + remainingActors: this.#actors.size, + elapsed, + }); + return true; + } else if (elapsed >= shutdownTimeout) { + this.log?.warn({ + msg: "shutdown timeout reached, forcing close", + remainingActors: this.#actors.size, + elapsed, + }); + return true; + } else { + // Log progress every 5 seconds + if (now - lastProgressLogTs >= progressLogInterval) { + this.log?.info({ + msg: "waiting for actors to stop", + remainingActors: this.#actors.size, + elapsed, + }); + lastProgressLogTs = now; + } + return false; + } + }; + + // Check immediately first + if (checkActors()) { + this.log?.debug({ + msg: "actors check completed immediately", + }); + resolve(); + return; + } + + this.log?.debug({ + msg: "starting actor wait interval", + checkInterval: shutdownCheckInterval, + }); + + const interval = setInterval(() => { + this.log?.debug({ + msg: "actor wait interval tick", + actorCount: this.#actors.size, + }); + if (checkActors()) { + this.log?.debug({ + msg: "actors check completed, clearing interval", + }); + clearInterval(interval); + resolve(); + } + }, shutdownCheckInterval); + }); + } + + // MARK: Networking + get pegboardEndpoint() { + return this.#config.pegboardEndpoint || this.#config.endpoint; + } + get pegboardUrl() { + const wsEndpoint = this.pegboardEndpoint + .replace("http://", "ws://") + .replace("https://", "wss://"); + + // Ensure the endpoint ends with /runners/connect + const baseUrl = wsEndpoint.endsWith("/") + ? wsEndpoint.slice(0, -1) + : wsEndpoint; + return `${baseUrl}/runners/connect?protocol_version=${PROTOCOL_VERSION}&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${encodeURIComponent(this.#runnerKey)}`; + } + + // MARK: Runner protocol + async #openPegboardWebSocket() { + const protocols = ["rivet"]; + if (this.config.token) + protocols.push(`rivet_token.${this.config.token}`); + + const WS = await importWebSocket(); + + // Assertion to clear previous WebSocket + if ( + this.#pegboardWebSocket && + (this.#pegboardWebSocket.readyState === WS.CONNECTING || + this.#pegboardWebSocket.readyState === WS.OPEN) + ) { + this.log?.error( + "found duplicate pegboardWebSocket, closing previous", + ); + this.#pegboardWebSocket.close(1000, "duplicate_websocket"); + } + + const ws = new WS(this.pegboardUrl, protocols) as any as WebSocket; + this.#pegboardWebSocket = ws; + + this.log?.info({ + msg: "connecting", + endpoint: this.pegboardEndpoint, + namespace: this.#config.namespace, + runnerKey: this.#runnerKey, + hasToken: !!this.config.token, + }); + + ws.addEventListener("open", () => { + if (this.#reconnectAttempt > 0) { + this.log?.info({ + msg: "runner reconnected", + namespace: this.#config.namespace, + runnerName: this.#config.runnerName, + reconnectAttempt: this.#reconnectAttempt, + }); + } else { + this.log?.debug({ + msg: "runner connected", + namespace: this.#config.namespace, + runnerName: this.#config.runnerName, + }); + } + + // Reset reconnect attempt counter on successful connection + this.#reconnectAttempt = 0; + + // Clear any pending reconnect timeout + if (this.#reconnectTimeout) { + clearTimeout(this.#reconnectTimeout); + this.#reconnectTimeout = undefined; + } + + // Clear any pending runner lost timeout since we're reconnecting + if (this.#runnerLostTimeout) { + clearTimeout(this.#runnerLostTimeout); + this.#runnerLostTimeout = undefined; + } + + // Send init message + const init: protocol.ToServerInit = { + name: this.#config.runnerName, + version: this.#config.version, + totalSlots: this.#config.totalSlots, + prepopulateActorNames: new Map( + Object.entries(this.#config.prepopulateActorNames).map( + ([name, data]) => [ + name, + { metadata: JSON.stringify(data.metadata) }, + ], + ), + ), + metadata: JSON.stringify(this.#config.metadata), + }; + + this.__sendToServer({ + tag: "ToServerInit", + val: init, + }); + + // Start command acknowledgment interval (5 minutes) + const ackInterval = 5 * 60 * 1000; // 5 minutes in milliseconds + const ackLoop = setInterval(() => { + try { + if (ws.readyState === 1) { + this.#sendCommandAcknowledgment(); + } else { + clearInterval(ackLoop); + this.log?.info({ + msg: "WebSocket not open, stopping ack loop", + }); + } + } catch (err) { + this.log?.error({ + msg: "error in command acknowledgment loop", + error: stringifyError(err), + }); + } + }, ackInterval); + this.#ackInterval = ackLoop; + }); + + ws.addEventListener("message", async (ev) => { + let buf: Uint8Array; + if (ev.data instanceof Blob) { + buf = new Uint8Array(await ev.data.arrayBuffer()); + } else if (Buffer.isBuffer(ev.data)) { + buf = new Uint8Array(ev.data); + } else { + throw new Error(`expected binary data, got ${typeof ev.data}`); + } + + await this.#injectLatency(); + + // Parse message + const message = protocol.decodeToClient(buf); + this.log?.debug({ + msg: "received runner message", + data: stringifyToClient(message), + }); + + // Handle message + if (message.tag === "ToClientInit") { + const init = message.val; + + if (this.runnerId !== init.runnerId) { + this.runnerId = init.runnerId; + + // Clear actors if runner id changed + this.#stopAllActors(); + } + + this.#protocolMetadata = init.metadata; + + this.log?.info({ + msg: "received init", + protocolMetadata: this.#protocolMetadata, + }); + + // Resend pending events + this.#processUnsentKvRequests(); + this.#resendUnacknowledgedEvents(); + this.#tunnel?.resendBufferedEvents(); + + this.#config.onConnected(); + } else if (message.tag === "ToClientCommands") { + const commands = message.val; + this.#handleCommands(commands); + } else if (message.tag === "ToClientAckEvents") { + this.#handleAckEvents(message.val); + } else if (message.tag === "ToClientKvResponse") { + const kvResponse = message.val; + this.#handleKvResponse(kvResponse); + } else if (message.tag === "ToClientTunnelMessage") { + this.#tunnel?.handleTunnelMessage(message.val).catch((err) => { + this.log?.error({ + msg: "error handling tunnel message", + error: stringifyError(err), + }); + }); + } else if (message.tag === "ToClientPing") { + this.__sendToServer({ + tag: "ToServerPong", + val: { + ts: message.val.ts, + }, + }); + } else { + unreachable(message); + } + }); + + ws.addEventListener("error", (ev) => { + this.log?.error({ + msg: `WebSocket error: ${stringifyError(ev.error)}`, + }); + + if (!this.#shutdown) { + this.#startRunnerLostTimeout(); + + // Attempt to reconnect if not stopped + this.#scheduleReconnect(); + } + }); + + ws.addEventListener("close", async (ev) => { + if (!this.#shutdown) { + const closeError = parseWebSocketCloseReason(ev.reason); + if ( + closeError?.group === "ws" && + closeError?.error === "eviction" + ) { + this.log?.info("runner websocket evicted"); + + this.#config.onDisconnected(ev.code, ev.reason); + + await this.shutdown(true); + } else { + this.log?.warn({ + msg: "runner disconnected", + code: ev.code, + reason: ev.reason.toString(), + closeError, + }); + + this.#config.onDisconnected(ev.code, ev.reason); + } + + // Clear ack interval on close + if (this.#ackInterval) { + clearInterval(this.#ackInterval); + this.#ackInterval = undefined; + } + + this.#startRunnerLostTimeout(); + + // Attempt to reconnect if not stopped + this.#scheduleReconnect(); + } else { + this.log?.info("websocket closed"); + + this.#config.onDisconnected(ev.code, ev.reason); + } + }); + } + + #startRunnerLostTimeout() { + // Start runner lost timeout if we have a threshold and are not shutting down + if ( + !this.#runnerLostTimeout && + this.#protocolMetadata && + this.#protocolMetadata.runnerLostThreshold > 0 + ) { + this.log?.info({ + msg: "starting runner lost timeout", + seconds: this.#protocolMetadata.runnerLostThreshold / 1000n, + }); + this.#runnerLostTimeout = setTimeout(() => { + try { + this.#handleLost(); + } catch (err) { + this.log?.error({ + msg: "error handling runner lost", + error: stringifyError(err), + }); + } + }, Number(this.#protocolMetadata.runnerLostThreshold)); + } + } + + #handleCommands(commands: protocol.ToClientCommands) { + this.log?.info({ + msg: "received commands", + commandCount: commands.length, + }); + + for (const commandWrapper of commands) { + if (commandWrapper.inner.tag === "CommandStartActor") { + // Spawn background promise + this.#handleCommandStartActor(commandWrapper).catch((err) => { + this.log?.error({ + msg: "error handling start actor command", + actorId: commandWrapper.checkpoint.actorId, + error: stringifyError(err), + }); + }); + + // NOTE: We don't do this for CommandStopActor because the actor will be removed by that call + // so we cant update the checkpoint + const actor = this.getActor( + commandWrapper.checkpoint.actorId, + commandWrapper.checkpoint.generation, + ); + if (actor) + actor.lastCommandIdx = commandWrapper.checkpoint.index; + } else if (commandWrapper.inner.tag === "CommandStopActor") { + // Spawn background promise + this.#handleCommandStopActor(commandWrapper).catch((err) => { + this.log?.error({ + msg: "error handling stop actor command", + actorId: commandWrapper.checkpoint.actorId, + error: stringifyError(err), + }); + }); + } else { + unreachable(commandWrapper.inner); + } + } + } + + #handleAckEvents(ack: protocol.ToClientAckEvents) { + const originalTotalEvents = Array.from(this.#actors).reduce( + (s, [_, actor]) => s + actor.eventHistory.length, + 0, + ); + + for (const [_, actor] of this.#actors) { + const checkpoint = ack.lastEventCheckpoints.find( + (x) => x.actorId == actor.actorId, + ); + + if (checkpoint) actor.handleAckEvents(checkpoint.index); + } + + const totalEvents = Array.from(this.#actors).reduce( + (s, [_, actor]) => s + actor.eventHistory.length, + 0, + ); + const prunedCount = originalTotalEvents - totalEvents; + + if (prunedCount > 0) { + this.log?.info({ + msg: "pruned acknowledged events", + prunedCount, + }); + } + + if (totalEvents <= EVENT_BACKLOG_WARN_THRESHOLD) { + this.#eventBacklogWarned = false; + } + } + + /** Track events to send to the server in case we need to resend it on disconnect. */ + #recordEvent(eventWrapper: protocol.EventWrapper) { + const actor = this.getActor(eventWrapper.checkpoint.actorId); + if (!actor) return; + + actor.recordEvent(eventWrapper); + + const totalEvents = Array.from(this.#actors).reduce( + (s, [_, actor]) => s + actor.eventHistory.length, + 0, + ); + + if ( + totalEvents > EVENT_BACKLOG_WARN_THRESHOLD && + !this.#eventBacklogWarned + ) { + this.#eventBacklogWarned = true; + this.log?.warn({ + msg: "unacknowledged event backlog exceeds threshold", + backlogSize: totalEvents, + threshold: EVENT_BACKLOG_WARN_THRESHOLD, + }); + } + } + + async #handleCommandStartActor(commandWrapper: protocol.CommandWrapper) { + // IMPORTANT: Make sure no async code runs before inserting #actors and + // calling addRequestToActor in order to prevent race conditions with + // subsequence commands + + if (!this.#tunnel) throw new Error("missing tunnel on actor start"); + + const startCommand = commandWrapper.inner + .val as protocol.CommandStartActor; + + const actorId = commandWrapper.checkpoint.actorId; + const generation = commandWrapper.checkpoint.generation; + const config = startCommand.config; + + const actorConfig: ActorConfig = { + name: config.name, + key: config.key, + createTs: config.createTs, + input: config.input ? new Uint8Array(config.input) : null, + }; + + const instance = new RunnerActor( + actorId, + generation, + actorConfig, + startCommand.hibernatingRequests, + ); + + const existingActor = this.#actors.get(actorId); + if (existingActor) { + this.log?.warn({ + msg: "replacing existing actor in actors map", + actorId, + existingGeneration: existingActor.generation, + newGeneration: generation, + existingPendingRequests: existingActor.pendingRequests.length, + }); + } + + this.#actors.set(actorId, instance); + + // NOTE: We have to populate the requestToActor map BEFORE running any + // async code in order for incoming tunnel messages to wait for + // instance.actorStartPromise before processing messages + // TODO: Where is this GC'd if something fails? + for (const hr of startCommand.hibernatingRequests) { + this.#tunnel.addRequestToActor(hr.gatewayId, hr.requestId, actorId); + } + + this.log?.info({ + msg: "created actor", + actors: this.#actors.size, + actorId, + name: config.name, + key: config.key, + generation, + hibernatingRequests: startCommand.hibernatingRequests.length, + }); + + this.#sendActorStateUpdate(actorId, generation, "running"); + + try { + // TODO: Add timeout to onActorStart + // Call onActorStart asynchronously and handle errors + this.log?.debug({ + msg: "calling onActorStart", + actorId, + generation, + }); + await this.#config.onActorStart(actorId, generation, actorConfig); + + instance.actorStartPromise.resolve(); + } catch (err) { + this.log?.error({ + msg: "error starting runner actor", + actorId, + err, + }); + + instance.actorStartPromise.reject(err); + + // TODO: Mark as crashed + // Send stopped state update if start failed + await this.forceStopActor(actorId, generation); + } + } + + async #handleCommandStopActor(commandWrapper: protocol.CommandWrapper) { + const stopCommand = commandWrapper.inner + .val as protocol.CommandStopActor; + + const actorId = commandWrapper.checkpoint.actorId; + const generation = commandWrapper.checkpoint.generation; + + await this.forceStopActor(actorId, generation); + } + + #sendActorIntent( + actorId: string, + generation: number, + intentType: "sleep" | "stop", + ) { + const actor = this.getActor(actorId, generation); + if (!actor) return; + + let actorIntent: protocol.ActorIntent; + + if (intentType === "sleep") { + actorIntent = { tag: "ActorIntentSleep", val: null }; + } else if (intentType === "stop") { + actorIntent = { + tag: "ActorIntentStop", + val: null, + }; + } else { + unreachable(intentType); + } + + const intentEvent: protocol.EventActorIntent = { + intent: actorIntent, + }; + + const eventWrapper: protocol.EventWrapper = { + checkpoint: { + actorId, + generation, + index: actor.nextEventIdx++, + }, + inner: { + tag: "EventActorIntent", + val: intentEvent, + }, + }; + + this.#recordEvent(eventWrapper); + + this.__sendToServer({ + tag: "ToServerEvents", + val: [eventWrapper], + }); + } + + #sendActorStateUpdate( + actorId: string, + generation: number, + stateType: "running" | "stopped", + ) { + const actor = this.getActor(actorId, generation); + if (!actor) return; + + let actorState: protocol.ActorState; + + if (stateType === "running") { + actorState = { tag: "ActorStateRunning", val: null }; + } else if (stateType === "stopped") { + actorState = { + tag: "ActorStateStopped", + val: { + code: actor.stopIntentSent || this.#draining + ? protocol.StopCode.Ok + : protocol.StopCode.Error, + message: null, + }, + }; + } else { + unreachable(stateType); + } + + const stateUpdateEvent: protocol.EventActorStateUpdate = { + state: actorState, + }; + + const eventWrapper: protocol.EventWrapper = { + checkpoint: { + actorId, + generation, + index: actor.nextEventIdx++, + }, + inner: { + tag: "EventActorStateUpdate", + val: stateUpdateEvent, + }, + }; + + this.#recordEvent(eventWrapper); + + this.__sendToServer({ + tag: "ToServerEvents", + val: [eventWrapper], + }); + } + + #sendCommandAcknowledgment() { + const lastCommandCheckpoints = []; + + for (const [_, actor] of this.#actors) { + if (actor.lastCommandIdx < 0) { + // No commands received yet, nothing to acknowledge + continue; + } + + lastCommandCheckpoints.push({ + actorId: actor.actorId, + generation: actor.generation, + index: actor.lastCommandIdx, + }); + } + + //this.#log?.log("Sending command acknowledgment", this.#lastCommandIdx); + + this.__sendToServer({ + tag: "ToServerAckCommands", + val: { + lastCommandCheckpoints, + }, + }); + } + + #handleKvResponse(response: protocol.ToClientKvResponse) { + const requestId = response.requestId; + const request = this.#kvRequests.get(requestId); + + if (!request) { + this.log?.error({ + msg: "received kv response for unknown request id", + requestId, + }); + return; + } + + this.#kvRequests.delete(requestId); + + if (response.data.tag === "KvErrorResponse") { + request.reject( + new Error(response.data.val.message || "Unknown KV error"), + ); + } else { + request.resolve(response.data.val); + } + } + + #parseGetResponseSimple( + response: protocol.KvGetResponse, + requestedKeys: Uint8Array[], + ): (Uint8Array | null)[] { + // Parse the response keys and values + const responseKeys: Uint8Array[] = []; + const responseValues: Uint8Array[] = []; + + for (const key of response.keys) { + responseKeys.push(new Uint8Array(key)); + } + + for (const value of response.values) { + responseValues.push(new Uint8Array(value)); + } + + // Map response back to requested key order + const result: (Uint8Array | null)[] = []; + for (const requestedKey of requestedKeys) { + let found = false; + for (let i = 0; i < responseKeys.length; i++) { + if (this.#keysEqual(requestedKey, responseKeys[i])) { + result.push(responseValues[i]); + found = true; + break; + } + } + if (!found) { + result.push(null); + } + } + + return result; + } + + #keysEqual(key1: Uint8Array, key2: Uint8Array): boolean { + if (key1.length !== key2.length) return false; + for (let i = 0; i < key1.length; i++) { + if (key1[i] !== key2[i]) return false; + } + return true; + } + + //#parseGetResponse(response: protocol.KvGetResponse) { + // const keys: string[] = []; + // const values: Uint8Array[] = []; + // const metadata: { version: Uint8Array; createTs: bigint }[] = []; + // + // for (const key of response.keys) { + // keys.push(new TextDecoder().decode(key)); + // } + // + // for (const value of response.values) { + // values.push(new Uint8Array(value)); + // } + // + // for (const meta of response.metadata) { + // metadata.push({ + // version: new Uint8Array(meta.version), + // createTs: meta.createTs, + // }); + // } + // + // return { keys, values, metadata }; + //} + + #parseListResponseSimple( + response: protocol.KvListResponse, + ): [Uint8Array, Uint8Array][] { + const result: [Uint8Array, Uint8Array][] = []; + + for (let i = 0; i < response.keys.length; i++) { + const key = response.keys[i]; + const value = response.values[i]; + + if (key && value) { + const keyBytes = new Uint8Array(key); + const valueBytes = new Uint8Array(value); + result.push([keyBytes, valueBytes]); + } + } + + return result; + } + + //#parseListResponse(response: protocol.KvListResponse) { + // const keys: string[] = []; + // const values: Uint8Array[] = []; + // const metadata: { version: Uint8Array; createTs: bigint }[] = []; + // + // for (const key of response.keys) { + // keys.push(new TextDecoder().decode(key)); + // } + // + // for (const value of response.values) { + // values.push(new Uint8Array(value)); + // } + // + // for (const meta of response.metadata) { + // metadata.push({ + // version: new Uint8Array(meta.version), + // createTs: meta.createTs, + // }); + // } + // + // return { keys, values, metadata }; + //} + + // MARK: KV Operations + async kvGet( + actorId: string, + keys: Uint8Array[], + ): Promise<(Uint8Array | null)[]> { + const kvKeys: protocol.KvKey[] = keys.map( + (key) => + key.buffer.slice( + key.byteOffset, + key.byteOffset + key.byteLength, + ) as ArrayBuffer, + ); + + const requestData: protocol.KvRequestData = { + tag: "KvGetRequest", + val: { keys: kvKeys }, + }; + + const response = await this.#sendKvRequest(actorId, requestData); + return this.#parseGetResponseSimple(response, keys); + } + + async kvListAll( + actorId: string, + options?: KvListOptions, + ): Promise<[Uint8Array, Uint8Array][]> { + const requestData: protocol.KvRequestData = { + tag: "KvListRequest", + val: { + query: { tag: "KvListAllQuery", val: null }, + reverse: options?.reverse || null, + limit: + options?.limit !== undefined ? BigInt(options.limit) : null, + }, + }; + + const response = await this.#sendKvRequest(actorId, requestData); + return this.#parseListResponseSimple(response); + } + + async kvListRange( + actorId: string, + start: Uint8Array, + end: Uint8Array, + exclusive?: boolean, + options?: KvListOptions, + ): Promise<[Uint8Array, Uint8Array][]> { + const startKey: protocol.KvKey = start.buffer.slice( + start.byteOffset, + start.byteOffset + start.byteLength, + ) as ArrayBuffer; + const endKey: protocol.KvKey = end.buffer.slice( + end.byteOffset, + end.byteOffset + end.byteLength, + ) as ArrayBuffer; + + const requestData: protocol.KvRequestData = { + tag: "KvListRequest", + val: { + query: { + tag: "KvListRangeQuery", + val: { + start: startKey, + end: endKey, + exclusive: exclusive || false, + }, + }, + reverse: options?.reverse || null, + limit: + options?.limit !== undefined ? BigInt(options.limit) : null, + }, + }; + + const response = await this.#sendKvRequest(actorId, requestData); + return this.#parseListResponseSimple(response); + } + + async kvListPrefix( + actorId: string, + prefix: Uint8Array, + options?: KvListOptions, + ): Promise<[Uint8Array, Uint8Array][]> { + const prefixKey: protocol.KvKey = prefix.buffer.slice( + prefix.byteOffset, + prefix.byteOffset + prefix.byteLength, + ) as ArrayBuffer; + + const requestData: protocol.KvRequestData = { + tag: "KvListRequest", + val: { + query: { + tag: "KvListPrefixQuery", + val: { key: prefixKey }, + }, + reverse: options?.reverse || null, + limit: + options?.limit !== undefined ? BigInt(options.limit) : null, + }, + }; + + const response = await this.#sendKvRequest(actorId, requestData); + return this.#parseListResponseSimple(response); + } + + async kvPut( + actorId: string, + entries: [Uint8Array, Uint8Array][], + ): Promise { + const keys: protocol.KvKey[] = entries.map( + ([key, _value]) => + key.buffer.slice( + key.byteOffset, + key.byteOffset + key.byteLength, + ) as ArrayBuffer, + ); + const values: protocol.KvValue[] = entries.map( + ([_key, value]) => + value.buffer.slice( + value.byteOffset, + value.byteOffset + value.byteLength, + ) as ArrayBuffer, + ); + + const requestData: protocol.KvRequestData = { + tag: "KvPutRequest", + val: { keys, values }, + }; + + await this.#sendKvRequest(actorId, requestData); + } + + async kvDelete(actorId: string, keys: Uint8Array[]): Promise { + const kvKeys: protocol.KvKey[] = keys.map( + (key) => + key.buffer.slice( + key.byteOffset, + key.byteOffset + key.byteLength, + ) as ArrayBuffer, + ); + + const requestData: protocol.KvRequestData = { + tag: "KvDeleteRequest", + val: { keys: kvKeys }, + }; + + await this.#sendKvRequest(actorId, requestData); + } + + async kvDeleteRange( + actorId: string, + start: Uint8Array, + end: Uint8Array, + ): Promise { + const startKey: protocol.KvKey = start.buffer.slice( + start.byteOffset, + start.byteOffset + start.byteLength, + ) as ArrayBuffer; + const endKey: protocol.KvKey = end.buffer.slice( + end.byteOffset, + end.byteOffset + end.byteLength, + ) as ArrayBuffer; + + const requestData: protocol.KvRequestData = { + tag: "KvDeleteRangeRequest", + val: { + start: startKey, + end: endKey, + }, + }; + + await this.#sendKvRequest(actorId, requestData); + } + + async kvDrop(actorId: string): Promise { + const requestData: protocol.KvRequestData = { + tag: "KvDropRequest", + val: null, + }; + + await this.#sendKvRequest(actorId, requestData); + } + + // MARK: Alarm Operations + setAlarm(actorId: string, alarmTs: number | null, generation?: number) { + const actor = this.getActor(actorId, generation); + if (!actor) return; + + const alarmEvent: protocol.EventActorSetAlarm = { + alarmTs: alarmTs !== null ? BigInt(alarmTs) : null, + }; + + const eventWrapper: protocol.EventWrapper = { + checkpoint: { + actorId, + generation: actor.generation, + index: actor.nextEventIdx++, + }, + inner: { + tag: "EventActorSetAlarm", + val: alarmEvent, + }, + }; + + this.#recordEvent(eventWrapper); + + this.__sendToServer({ + tag: "ToServerEvents", + val: [eventWrapper], + }); + } + + clearAlarm(actorId: string, generation?: number) { + this.setAlarm(actorId, null, generation); + } + + #sendKvRequest( + actorId: string, + requestData: protocol.KvRequestData, + ): Promise { + return new Promise((resolve, reject) => { + const requestId = this.#nextKvRequestId++; + + // Store the request + const requestEntry = { + actorId, + data: requestData, + resolve, + reject, + sent: false, + timestamp: Date.now(), + }; + + this.#kvRequests.set(requestId, requestEntry); + + if (this.getPegboardWebSocketIfReady()) { + // Send immediately + this.#sendSingleKvRequest(requestId); + } + }); + } + + #sendSingleKvRequest(requestId: number) { + const request = this.#kvRequests.get(requestId); + if (!request || request.sent) return; + + try { + const kvRequest: protocol.ToServerKvRequest = { + actorId: request.actorId, + requestId, + data: request.data, + }; + + this.__sendToServer({ + tag: "ToServerKvRequest", + val: kvRequest, + }); + + // Mark as sent and update timestamp + request.sent = true; + request.timestamp = Date.now(); + } catch (error) { + this.#kvRequests.delete(requestId); + request.reject(error); + } + } + + #processUnsentKvRequests() { + if (!this.getPegboardWebSocketIfReady()) { + return; + } + + let processedCount = 0; + for (const [requestId, request] of this.#kvRequests.entries()) { + if (!request.sent) { + this.#sendSingleKvRequest(requestId); + processedCount++; + } + } + + if (processedCount > 0) { + //this.#log?.log(`Processed ${processedCount} queued KV requests`); + } + } + + /** Resolves after the configured debug latency, or immediately if none. */ + #injectLatency(): Promise { + const ms = this.#config.debugLatencyMs; + if (!ms) return Promise.resolve(); + return new Promise((resolve) => setTimeout(resolve, ms)); + } + + /** Asserts WebSocket exists and is ready. */ + getPegboardWebSocketIfReady(): WebSocket | undefined { + if ( + !!this.#pegboardWebSocket && + this.#pegboardWebSocket.readyState === 1 + ) { + return this.#pegboardWebSocket; + } else { + return undefined; + } + } + + __sendToServer(message: protocol.ToServer) { + this.log?.debug({ + msg: "sending runner message", + data: stringifyToServer(message), + }); + + const encoded = protocol.encodeToServer(message); + + // Normally synchronous. When debugLatencyMs is set, the send is + // deferred but message order is preserved. + this.#injectLatency().then(() => { + const pegboardWebSocket = this.getPegboardWebSocketIfReady(); + if (pegboardWebSocket) { + pegboardWebSocket.send(encoded); + } else { + this.log?.error({ + msg: "WebSocket not available or not open for sending data", + }); + } + }); + } + + sendHibernatableWebSocketMessageAck( + gatewayId: ArrayBuffer, + requestId: ArrayBuffer, + index: number, + ) { + if (!this.#tunnel) + throw new Error("missing tunnel to send message ack"); + this.#tunnel.sendHibernatableWebSocketMessageAck( + gatewayId, + requestId, + index, + ); + } + + /** + * Restores hibernatable WebSocket connections for an actor. + * + * This method should be called at the end of `onActorStart` after the + * actor instance is fully initialized. + * + * This method will: + * - Restore all provided hibernatable WebSocket connections + * - Attach event listeners to the restored WebSockets + * - Close any WebSocket connections that failed to restore + * + * The provided metadata list should include all hibernatable WebSockets + * that were persisted for this actor. The gateway will automatically + * close any connections that are not restored (i.e., not included in + * this list). + * + * **Important:** This method must be called after `onActorStart` completes + * and before marking the actor as "ready" to ensure all hibernatable + * connections are fully restored. + * + * @param actorId - The ID of the actor to restore connections for + * @param metaEntries - Array of hibernatable WebSocket metadata to restore + */ + async restoreHibernatingRequests( + actorId: string, + metaEntries: HibernatingWebSocketMetadata[], + ) { + if (!this.#tunnel) + throw new Error("missing tunnel to restore hibernating requests"); + await this.#tunnel.restoreHibernatingRequests(actorId, metaEntries); + } + + getServerlessInitPacket(): string | undefined { + if (!this.runnerId) return undefined; + + const data = protocol.encodeToServerlessServer({ + tag: "ToServerlessServerInit", + val: { + runnerId: this.runnerId, + runnerProtocolVersion: PROTOCOL_VERSION, + }, + }); + + // Embed version + const buffer = Buffer.alloc(data.length + 2); + buffer.writeUInt16LE(PROTOCOL_VERSION, 0); + Buffer.from(data).copy(buffer, 2); + + return buffer.toString("base64"); + } + + #scheduleReconnect() { + if (this.#shutdown) { + this.log?.debug({ + msg: "Runner is shut down, not attempting reconnect", + }); + return; + } + + const delay = calculateBackoff(this.#reconnectAttempt, { + initialDelay: 1000, + maxDelay: 30000, + multiplier: 2, + jitter: true, + }); + + this.log?.debug({ + msg: `Scheduling reconnect attempt ${this.#reconnectAttempt + 1} in ${delay}ms`, + }); + + if (this.#reconnectTimeout) { + this.log?.info( + "clearing previous reconnect timeout in schedule reconnect", + ); + clearTimeout(this.#reconnectTimeout); + } + + this.#reconnectTimeout = setTimeout(() => { + if (!this.#shutdown) { + this.#reconnectAttempt++; + this.log?.debug({ + msg: `Attempting to reconnect (attempt ${this.#reconnectAttempt})...`, + }); + this.#openPegboardWebSocket().catch((err) => { + this.log?.error({ + msg: "error during websocket reconnection", + error: stringifyError(err), + }); + }); + } + }, delay); + } + + #resendUnacknowledgedEvents() { + const eventsToResend = []; + + for (const [_, actor] of this.#actors) { + eventsToResend.push(...actor.eventHistory); + } + + if (eventsToResend.length === 0) return; + + this.log?.info({ + msg: "resending unacknowledged events", + count: eventsToResend.length, + }); + + // Resend events in batches + this.__sendToServer({ + tag: "ToServerEvents", + val: eventsToResend, + }); + } + + #cleanupOldKvRequests() { + const thirtySecondsAgo = Date.now() - KV_EXPIRE; + const toDelete: number[] = []; + + for (const [requestId, request] of this.#kvRequests.entries()) { + if (request.timestamp < thirtySecondsAgo) { + request.reject( + new Error( + "KV request timed out waiting for WebSocket connection", + ), + ); + toDelete.push(requestId); + } + } + + for (const requestId of toDelete) { + this.#kvRequests.delete(requestId); + } + + if (toDelete.length > 0) { + //this.#log?.log(`Cleaned up ${toDelete.length} expired KV requests`); + } + } + + getProtocolMetadata(): protocol.ProtocolMetadata | undefined { + return this.#protocolMetadata; + } +} diff --git a/rivetkit-typescript/packages/engine-runner/src/stringify.ts b/rivetkit-typescript/packages/engine-runner/src/stringify.ts new file mode 100644 index 0000000000..f32226f511 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/stringify.ts @@ -0,0 +1,357 @@ +import type * as protocol from "@rivetkit/engine-runner-protocol"; +import { idToStr } from "./utils"; + +/** + * Helper function to stringify ArrayBuffer for logging + */ +function stringifyArrayBuffer(buffer: ArrayBuffer): string { + return `ArrayBuffer(${buffer.byteLength})`; +} + +/** + * Helper function to stringify bigint for logging + */ +function stringifyBigInt(value: bigint): string { + return `${value}n`; +} + +/** + * Helper function to stringify Map for logging + */ +function stringifyMap(map: ReadonlyMap): string { + const entries = Array.from(map.entries()) + .map(([k, v]) => `"${k}": "${v}"`) + .join(", "); + return `Map(${map.size}){${entries}}`; +} + +/** + * Helper function to stringify MessageId for logging + */ +function stringifyMessageId(messageId: protocol.MessageId): string { + return `MessageId{gatewayId: ${idToStr(messageId.gatewayId)}, requestId: ${idToStr(messageId.requestId)}, messageIndex: ${messageId.messageIndex}}`; +} + +/** + * Stringify ToServerTunnelMessageKind for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyToServerTunnelMessageKind( + kind: protocol.ToServerTunnelMessageKind, +): string { + switch (kind.tag) { + case "ToServerResponseStart": { + const { status, headers, body, stream } = kind.val; + const bodyStr = body === null ? "null" : stringifyArrayBuffer(body); + return `ToServerResponseStart{status: ${status}, headers: ${stringifyMap(headers)}, body: ${bodyStr}, stream: ${stream}}`; + } + case "ToServerResponseChunk": { + const { body, finish } = kind.val; + return `ToServerResponseChunk{body: ${stringifyArrayBuffer(body)}, finish: ${finish}}`; + } + case "ToServerResponseAbort": + return "ToServerResponseAbort"; + case "ToServerWebSocketOpen": { + const { canHibernate } = kind.val; + return `ToServerWebSocketOpen{canHibernate: ${canHibernate}}`; + } + case "ToServerWebSocketMessage": { + const { data, binary } = kind.val; + return `ToServerWebSocketMessage{data: ${stringifyArrayBuffer(data)}, binary: ${binary}}`; + } + case "ToServerWebSocketMessageAck": { + const { index } = kind.val; + return `ToServerWebSocketMessageAck{index: ${index}}`; + } + case "ToServerWebSocketClose": { + const { code, reason, hibernate } = kind.val; + const codeStr = code === null ? "null" : code.toString(); + const reasonStr = reason === null ? "null" : `"${reason}"`; + return `ToServerWebSocketClose{code: ${codeStr}, reason: ${reasonStr}, hibernate: ${hibernate}}`; + } + } +} + +/** + * Stringify ToClientTunnelMessageKind for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyToClientTunnelMessageKind( + kind: protocol.ToClientTunnelMessageKind, +): string { + switch (kind.tag) { + case "ToClientRequestStart": { + const { actorId, method, path, headers, body, stream } = kind.val; + const bodyStr = body === null ? "null" : stringifyArrayBuffer(body); + return `ToClientRequestStart{actorId: "${actorId}", method: "${method}", path: "${path}", headers: ${stringifyMap(headers)}, body: ${bodyStr}, stream: ${stream}}`; + } + case "ToClientRequestChunk": { + const { body, finish } = kind.val; + return `ToClientRequestChunk{body: ${stringifyArrayBuffer(body)}, finish: ${finish}}`; + } + case "ToClientRequestAbort": + return "ToClientRequestAbort"; + case "ToClientWebSocketOpen": { + const { actorId, path, headers } = kind.val; + return `ToClientWebSocketOpen{actorId: "${actorId}", path: "${path}", headers: ${stringifyMap(headers)}}`; + } + case "ToClientWebSocketMessage": { + const { data, binary } = kind.val; + return `ToClientWebSocketMessage{data: ${stringifyArrayBuffer(data)}, binary: ${binary}}`; + } + case "ToClientWebSocketClose": { + const { code, reason } = kind.val; + const codeStr = code === null ? "null" : code.toString(); + const reasonStr = reason === null ? "null" : `"${reason}"`; + return `ToClientWebSocketClose{code: ${codeStr}, reason: ${reasonStr}}`; + } + } +} + +/** + * Stringify Command for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyCommand(command: protocol.Command): string { + switch (command.tag) { + case "CommandStartActor": { + const { config, hibernatingRequests } = command.val; + const keyStr = config.key === null ? "null" : `"${config.key}"`; + const inputStr = + config.input === null + ? "null" + : stringifyArrayBuffer(config.input); + const hibernatingRequestsStr = + hibernatingRequests.length > 0 + ? `[${hibernatingRequests.map((hr) => `{gatewayId: ${idToStr(hr.gatewayId)}, requestId: ${idToStr(hr.requestId)}}`).join(", ")}]` + : "[]"; + return `CommandStartActor{config: {name: "${config.name}", key: ${keyStr}, createTs: ${stringifyBigInt(config.createTs)}, input: ${inputStr}}, hibernatingRequests: ${hibernatingRequestsStr}}`; + } + case "CommandStopActor": { + return `CommandStopActor`; + } + } +} + +/** + * Stringify CommandWrapper for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyCommandWrapper( + wrapper: protocol.CommandWrapper, +): string { + return `CommandWrapper{actorId: "${wrapper.checkpoint.actorId}", generation: "${wrapper.checkpoint.generation}", index: ${stringifyBigInt(wrapper.checkpoint.index)}, inner: ${stringifyCommand(wrapper.inner)}}`; +} + +/** + * Stringify Event for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyEvent(event: protocol.Event): string { + switch (event.tag) { + case "EventActorIntent": { + const { intent } = event.val; + const intentStr = + intent.tag === "ActorIntentSleep" + ? "Sleep" + : intent.tag === "ActorIntentStop" + ? "Stop" + : "Unknown"; + return `EventActorIntent{intent: ${intentStr}}`; + } + case "EventActorStateUpdate": { + const { state } = event.val; + let stateStr: string; + if (state.tag === "ActorStateRunning") { + stateStr = "Running"; + } else if (state.tag === "ActorStateStopped") { + const { code, message } = state.val; + const messageStr = message === null ? "null" : `"${message}"`; + stateStr = `Stopped{code: ${code}, message: ${messageStr}}`; + } else { + stateStr = "Unknown"; + } + return `EventActorStateUpdate{state: ${stateStr}}`; + } + case "EventActorSetAlarm": { + const { alarmTs } = event.val; + const alarmTsStr = + alarmTs === null ? "null" : stringifyBigInt(alarmTs); + return `EventActorSetAlarm{alarmTs: ${alarmTsStr}}`; + } + } +} + +/** + * Stringify EventWrapper for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyEventWrapper(wrapper: protocol.EventWrapper): string { + return `EventWrapper{actorId: ${wrapper.checkpoint.actorId}, generation: "${wrapper.checkpoint.generation}", index: ${stringifyBigInt(wrapper.checkpoint.index)}, inner: ${stringifyEvent(wrapper.inner)}}`; +} + +/** + * Stringify ToServer for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyToServer(message: protocol.ToServer): string { + switch (message.tag) { + case "ToServerInit": { + const { + name, + version, + totalSlots, + prepopulateActorNames, + metadata, + } = message.val; + const prepopulateActorNamesStr = + prepopulateActorNames === null + ? "null" + : `Map(${prepopulateActorNames.size})`; + const metadataStr = metadata === null ? "null" : `"${metadata}"`; + return `ToServerInit{name: "${name}", version: ${version}, totalSlots: ${totalSlots}, prepopulateActorNames: ${prepopulateActorNamesStr}, metadata: ${metadataStr}}`; + } + case "ToServerEvents": { + const events = message.val; + return `ToServerEvents{count: ${events.length}, events: [${events.map((e) => stringifyEventWrapper(e)).join(", ")}]}`; + } + case "ToServerAckCommands": { + const { lastCommandCheckpoints } = message.val; + const checkpointsStr = + lastCommandCheckpoints.length > 0 + ? `[${lastCommandCheckpoints.map((cp) => `{actorId: "${cp.actorId}", index: ${stringifyBigInt(cp.index)}}`).join(", ")}]` + : "[]"; + return `ToServerAckCommands{lastCommandCheckpoints: ${checkpointsStr}}`; + } + case "ToServerStopping": + return "ToServerStopping"; + case "ToServerPong": { + const { ts } = message.val; + return `ToServerPong{ts: ${stringifyBigInt(ts)}}`; + } + case "ToServerKvRequest": { + const { actorId, requestId, data } = message.val; + const dataStr = stringifyKvRequestData(data); + return `ToServerKvRequest{actorId: "${actorId}", requestId: ${requestId}, data: ${dataStr}}`; + } + case "ToServerTunnelMessage": { + const { messageId, messageKind } = message.val; + return `ToServerTunnelMessage{messageId: ${stringifyMessageId(messageId)}, messageKind: ${stringifyToServerTunnelMessageKind(messageKind)}}`; + } + } +} + +/** + * Stringify ToClient for logging + * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified + */ +export function stringifyToClient(message: protocol.ToClient): string { + switch (message.tag) { + case "ToClientInit": { + const { runnerId, metadata } = message.val; + const metadataStr = `{runnerLostThreshold: ${stringifyBigInt(metadata.runnerLostThreshold)}}`; + return `ToClientInit{runnerId: "${runnerId}", metadata: ${metadataStr}}`; + } + case "ToClientPing": { + const { ts } = message.val; + return `ToClientPing{ts: ${stringifyBigInt(ts)}}`; + } + case "ToClientCommands": { + const commands = message.val; + return `ToClientCommands{count: ${commands.length}, commands: [${commands.map((c) => stringifyCommandWrapper(c)).join(", ")}]}`; + } + case "ToClientAckEvents": { + const { lastEventCheckpoints } = message.val; + const checkpointsStr = + lastEventCheckpoints.length > 0 + ? `[${lastEventCheckpoints.map((cp) => `{actorId: "${cp.actorId}", index: ${stringifyBigInt(cp.index)}}`).join(", ")}]` + : "[]"; + return `ToClientAckEvents{lastEventCheckpoints: ${checkpointsStr}}`; + } + case "ToClientKvResponse": { + const { requestId, data } = message.val; + const dataStr = stringifyKvResponseData(data); + return `ToClientKvResponse{requestId: ${requestId}, data: ${dataStr}}`; + } + case "ToClientTunnelMessage": { + const { messageId, messageKind } = message.val; + return `ToClientTunnelMessage{messageId: ${stringifyMessageId(messageId)}, messageKind: ${stringifyToClientTunnelMessageKind(messageKind)}}`; + } + } +} + +/** + * Stringify KvRequestData for logging + */ +function stringifyKvRequestData(data: protocol.KvRequestData): string { + switch (data.tag) { + case "KvGetRequest": { + const { keys } = data.val; + return `KvGetRequest{keys: ${keys.length}}`; + } + case "KvListRequest": { + const { query, reverse, limit } = data.val; + const reverseStr = reverse === null ? "null" : reverse.toString(); + const limitStr = limit === null ? "null" : stringifyBigInt(limit); + return `KvListRequest{query: ${stringifyKvListQuery(query)}, reverse: ${reverseStr}, limit: ${limitStr}}`; + } + case "KvPutRequest": { + const { keys, values } = data.val; + return `KvPutRequest{keys: ${keys.length}, values: ${values.length}}`; + } + case "KvDeleteRequest": { + const { keys } = data.val; + return `KvDeleteRequest{keys: ${keys.length}}`; + } + case "KvDeleteRangeRequest": { + const { start, end } = data.val; + return `KvDeleteRangeRequest{start: ${stringifyArrayBuffer(start)}, end: ${stringifyArrayBuffer(end)}}`; + } + case "KvDropRequest": + return "KvDropRequest"; + } +} + +/** + * Stringify KvListQuery for logging + */ +function stringifyKvListQuery(query: protocol.KvListQuery): string { + switch (query.tag) { + case "KvListAllQuery": + return "KvListAllQuery"; + case "KvListRangeQuery": { + const { start, end, exclusive } = query.val; + return `KvListRangeQuery{start: ${stringifyArrayBuffer(start)}, end: ${stringifyArrayBuffer(end)}, exclusive: ${exclusive}}`; + } + case "KvListPrefixQuery": { + const { key } = query.val; + return `KvListPrefixQuery{key: ${stringifyArrayBuffer(key)}}`; + } + } +} + +/** + * Stringify KvResponseData for logging + */ +function stringifyKvResponseData(data: protocol.KvResponseData): string { + switch (data.tag) { + case "KvErrorResponse": { + const { message } = data.val; + return `KvErrorResponse{message: "${message}"}`; + } + case "KvGetResponse": { + const { keys, values, metadata } = data.val; + return `KvGetResponse{keys: ${keys.length}, values: ${values.length}, metadata: ${metadata.length}}`; + } + case "KvListResponse": { + const { keys, values, metadata } = data.val; + return `KvListResponse{keys: ${keys.length}, values: ${values.length}, metadata: ${metadata.length}}`; + } + case "KvPutResponse": + return "KvPutResponse"; + case "KvDeleteResponse": + return "KvDeleteResponse"; + case "KvDropResponse": + return "KvDropResponse"; + } +} diff --git a/rivetkit-typescript/packages/engine-runner/src/tunnel.ts b/rivetkit-typescript/packages/engine-runner/src/tunnel.ts new file mode 100644 index 0000000000..ad2b650e74 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/tunnel.ts @@ -0,0 +1,1168 @@ +import type * as protocol from "@rivetkit/engine-runner-protocol"; +import type { + GatewayId, + MessageId, + RequestId, +} from "@rivetkit/engine-runner-protocol"; +import type { Logger } from "pino"; +import { type Runner, type RunnerActor, RunnerShutdownError } from "./mod"; +import { + stringifyToClientTunnelMessageKind, + stringifyToServerTunnelMessageKind, +} from "./stringify"; +import { arraysEqual, idToStr, MAX_PAYLOAD_SIZE, stringifyError, unreachable } from "./utils"; +import { + HIBERNATABLE_SYMBOL, + WebSocketTunnelAdapter, +} from "./websocket-tunnel-adapter"; + +export interface PendingRequest { + resolve: (response: Response) => void; + reject: (error: Error) => void; + streamController?: ReadableStreamDefaultController; + actorId?: string; + gatewayId?: GatewayId; + requestId?: RequestId; + clientMessageIndex: number; +} + +export interface HibernatingWebSocketMetadata { + gatewayId: GatewayId; + requestId: RequestId; + clientMessageIndex: number; + serverMessageIndex: number; + + path: string; + headers: Record; +} + +export class Tunnel { + #runner: Runner; + + /** Maps request IDs to actor IDs for lookup */ + #requestToActor: Array<{ + gatewayId: GatewayId; + requestId: RequestId; + actorId: string; + }> = []; + + /** Buffer for messages when not connected */ + #bufferedMessages: Array<{ + gatewayId: GatewayId; + requestId: RequestId; + messageKind: protocol.ToServerTunnelMessageKind; + }> = []; + + get log(): Logger | undefined { + return this.#runner.log; + } + + constructor(runner: Runner) { + this.#runner = runner; + } + + start(): void { + // No-op - kept for compatibility + } + + resendBufferedEvents(): void { + if (this.#bufferedMessages.length === 0) { + return; + } + + this.log?.info({ + msg: "resending buffered tunnel messages", + count: this.#bufferedMessages.length, + }); + + const messages = this.#bufferedMessages; + this.#bufferedMessages = []; + + for (const { gatewayId, requestId, messageKind } of messages) { + this.#sendMessage(gatewayId, requestId, messageKind); + } + } + + shutdown() { + // NOTE: Pegboard WS already closed at this point, cannot send + // anything. All teardown logic is handled by pegboard-runner. + + // Reject all pending requests and close all WebSockets for all actors + // RunnerShutdownError will be explicitly ignored + for (const [_actorId, actor] of this.#runner.actors) { + // Reject all pending requests for this actor + for (const entry of actor.pendingRequests) { + entry.request.reject(new RunnerShutdownError()); + } + actor.pendingRequests = []; + + // Close all WebSockets for this actor + // The WebSocket close event with retry is automatically sent when the + // runner WS closes, so we only need to notify the client that the WS + // closed: + // https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157 + for (const entry of actor.webSockets) { + // Only close non-hibernatable websockets to prevent sending + // unnecessary close messages for websockets that will be hibernated + if (!entry.ws[HIBERNATABLE_SYMBOL]) { + entry.ws._closeWithoutCallback(1000, "ws.tunnel_shutdown"); + } + } + actor.webSockets = []; + } + + // Clear the request-to-actor mapping + this.#requestToActor = []; + } + + async restoreHibernatingRequests( + actorId: string, + metaEntries: HibernatingWebSocketMetadata[], + ) { + const actor = this.#runner.getActor(actorId); + if (!actor) { + throw new Error( + `Actor ${actorId} not found for restoring hibernating requests`, + ); + } + + if (actor.hibernationRestored) { + throw new Error( + `Actor ${actorId} already restored hibernating requests`, + ); + } + + this.log?.debug({ + msg: "restoring hibernating requests", + actorId, + requests: actor.hibernatingRequests.length, + }); + + // Track all background operations + const backgroundOperations: Promise[] = []; + + // Process connected WebSockets + let connectedButNotLoadedCount = 0; + let restoredCount = 0; + for (const { gatewayId, requestId } of actor.hibernatingRequests) { + const requestIdStr = idToStr(requestId); + const meta = metaEntries.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + + if (!meta) { + // Connected but not loaded (not persisted) - close it + // + // This may happen if the metadata was not successfully persisted + this.log?.warn({ + msg: "closing websocket that is not persisted", + requestId: requestIdStr, + }); + + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketClose", + val: { + code: 1000, + reason: "ws.meta_not_found_during_restore", + hibernate: false, + }, + }); + + connectedButNotLoadedCount++; + } else { + // Both connected and persisted - restore it + const request = buildRequestForWebSocket( + meta.path, + meta.headers, + ); + + // This will call `runner.config.websocket` under the hood to + // attach the event listeners to the WebSocket. + // Track this operation to ensure it completes + const restoreOperation = this.#createWebSocket( + actorId, + gatewayId, + requestId, + requestIdStr, + meta.serverMessageIndex, + true, + true, + request, + meta.path, + meta.headers, + false, + ) + .then(() => { + // Create a PendingRequest entry to track the message index + const actor = this.#runner.getActor(actorId); + if (actor) { + actor.createPendingRequest( + gatewayId, + requestId, + meta.clientMessageIndex, + ); + } + + this.log?.info({ + msg: "connection successfully restored", + actorId, + requestId: requestIdStr, + }); + }) + .catch((err) => { + this.log?.error({ + msg: "error creating websocket during restore", + requestId: requestIdStr, + error: stringifyError(err), + }); + + // Close the WebSocket on error + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketClose", + val: { + code: 1011, + reason: "ws.restore_error", + hibernate: false, + }, + }); + }); + + backgroundOperations.push(restoreOperation); + restoredCount++; + } + } + + // Process loaded but not connected (stale) - remove them + let loadedButNotConnectedCount = 0; + for (const meta of metaEntries) { + const requestIdStr = idToStr(meta.requestId); + const isConnected = actor.hibernatingRequests.some( + (req) => + arraysEqual(req.gatewayId, meta.gatewayId) && + arraysEqual(req.requestId, meta.requestId), + ); + if (!isConnected) { + this.log?.warn({ + msg: "removing stale persisted websocket", + requestId: requestIdStr, + }); + + const request = buildRequestForWebSocket( + meta.path, + meta.headers, + ); + + // Create adapter to register user's event listeners. + // Pass engineAlreadyClosed=true so close callback won't send tunnel message. + // Track this operation to ensure it completes + const cleanupOperation = this.#createWebSocket( + actorId, + meta.gatewayId, + meta.requestId, + requestIdStr, + meta.serverMessageIndex, + true, + true, + request, + meta.path, + meta.headers, + true, + ) + .then((adapter) => { + // Close the adapter normally - this will fire user's close event handler + // (which should clean up persistence) and trigger the close callback + // (which will clean up maps but skip sending tunnel message) + adapter.close(1000, "ws.stale_metadata"); + }) + .catch((err) => { + this.log?.error({ + msg: "error creating stale websocket during restore", + requestId: requestIdStr, + error: stringifyError(err), + }); + }); + + backgroundOperations.push(cleanupOperation); + loadedButNotConnectedCount++; + } + } + + // Wait for all background operations to complete before finishing + await Promise.allSettled(backgroundOperations); + + // Mark restoration as complete + actor.hibernationRestored = true; + + this.log?.info({ + msg: "restored hibernatable websockets", + actorId, + restoredCount, + connectedButNotLoadedCount, + loadedButNotConnectedCount, + }); + } + + /** + * Called from WebSocketOpen message and when restoring hibernatable WebSockets. + * + * engineAlreadyClosed will be true if this is only being called to trigger + * the close callback and not to send a close message to the server. This + * is used specifically to clean up zombie WebSocket connections. + */ + async #createWebSocket( + actorId: string, + gatewayId: GatewayId, + requestId: RequestId, + requestIdStr: string, + serverMessageIndex: number, + isHibernatable: boolean, + isRestoringHibernatable: boolean, + request: Request, + path: string, + headers: Record, + engineAlreadyClosed: boolean, + ): Promise { + this.log?.debug({ + msg: "createWebSocket creating adapter", + actorId, + requestIdStr, + isHibernatable, + path, + }); + // Create WebSocket adapter + const adapter = new WebSocketTunnelAdapter( + this, + actorId, + requestIdStr, + serverMessageIndex, + isHibernatable, + isRestoringHibernatable, + request, + (data: ArrayBuffer | string, isBinary: boolean) => { + // Send message through tunnel + const dataBuffer = + typeof data === "string" + ? (new TextEncoder().encode(data).buffer as ArrayBuffer) + : data; + + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketMessage", + val: { + data: dataBuffer, + binary: isBinary, + }, + }); + }, + (code?: number, reason?: string) => { + // Send close through tunnel if engine doesn't already know it's closed + if (!engineAlreadyClosed) { + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketClose", + val: { + code: code || null, + reason: reason || null, + hibernate: false, + }, + }); + } + + // Clean up actor tracking + const actor = this.#runner.getActor(actorId); + if (actor) { + actor.deleteWebSocket(gatewayId, requestId); + actor.deletePendingRequest(gatewayId, requestId); + } + + // Clean up request-to-actor mapping + this.#removeRequestToActor(gatewayId, requestId); + }, + ); + + // Get actor and add websocket to it + const actor = this.#runner.getActor(actorId); + if (!actor) { + throw new Error(`Actor ${actorId} not found`); + } + + actor.setWebSocket(gatewayId, requestId, adapter); + this.addRequestToActor(gatewayId, requestId, actorId); + + // Call WebSocket handler. This handler will add event listeners + // for `open`, etc. Pass the VirtualWebSocket (not the adapter) to the actor. + await this.#runner.config.websocket( + this.#runner, + actorId, + adapter.websocket, + gatewayId, + requestId, + request, + path, + headers, + isHibernatable, + isRestoringHibernatable, + ); + + return adapter; + } + + addRequestToActor( + gatewayId: GatewayId, + requestId: RequestId, + actorId: string, + ) { + this.#requestToActor.push({ gatewayId, requestId, actorId }); + } + + #removeRequestToActor(gatewayId: GatewayId, requestId: RequestId) { + const index = this.#requestToActor.findIndex( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + if (index !== -1) { + this.#requestToActor.splice(index, 1); + } + } + + getRequestActor( + gatewayId: GatewayId, + requestId: RequestId, + ): RunnerActor | undefined { + const entry = this.#requestToActor.find( + (entry) => + arraysEqual(entry.gatewayId, gatewayId) && + arraysEqual(entry.requestId, requestId), + ); + + if (!entry) { + this.log?.warn({ + msg: "missing requestToActor entry", + requestId: idToStr(requestId), + }); + return undefined; + } + + const actor = this.#runner.getActor(entry.actorId); + if (!actor) { + this.log?.warn({ + msg: "missing actor for requestToActor lookup", + requestId: idToStr(requestId), + actorId: entry.actorId, + }); + return undefined; + } + + return actor; + } + + async getAndWaitForRequestActor( + gatewayId: GatewayId, + requestId: RequestId, + ): Promise { + const actor = this.getRequestActor(gatewayId, requestId); + if (!actor) return; + await actor.actorStartPromise.promise; + return actor; + } + + #sendMessage( + gatewayId: GatewayId, + requestId: RequestId, + messageKind: protocol.ToServerTunnelMessageKind, + ) { + // Buffer message if not connected + if (!this.#runner.getPegboardWebSocketIfReady()) { + this.log?.debug({ + msg: "buffering tunnel message, socket not connected to engine", + requestId: idToStr(requestId), + message: stringifyToServerTunnelMessageKind(messageKind), + }); + this.#bufferedMessages.push({ gatewayId, requestId, messageKind }); + return; + } + + // Get or initialize message index for this request + // + // We don't have to wait for the actor to start since we're not calling + // any callbacks on the actor + const gatewayIdStr = idToStr(gatewayId); + const requestIdStr = idToStr(requestId); + const actor = this.getRequestActor(gatewayId, requestId); + if (!actor) { + this.log?.warn({ + msg: "cannot send tunnel message, actor not found", + gatewayId: gatewayIdStr, + requestId: requestIdStr, + }); + return; + } + + // Get message index from pending request + let clientMessageIndex: number; + const pending = actor.getPendingRequest(gatewayId, requestId); + if (pending) { + clientMessageIndex = pending.clientMessageIndex; + pending.clientMessageIndex++; + } else { + // No pending request + this.log?.warn({ + msg: "missing pending request for send message, defaulting to message index 0", + gatewayId: gatewayIdStr, + requestId: requestIdStr, + }); + clientMessageIndex = 0; + } + + // Build message ID from gatewayId + requestId + messageIndex + const messageId: protocol.MessageId = { + gatewayId, + requestId, + messageIndex: clientMessageIndex, + }; + const messageIdStr = `${idToStr(messageId.gatewayId)}-${idToStr(messageId.requestId)}-${messageId.messageIndex}`; + + this.log?.debug({ + msg: "sending tunnel msg", + messageId: messageIdStr, + gatewayId: gatewayIdStr, + requestId: requestIdStr, + messageIndex: clientMessageIndex, + message: stringifyToServerTunnelMessageKind(messageKind), + }); + + // Send message + const message: protocol.ToServer = { + tag: "ToServerTunnelMessage", + val: { + messageId, + messageKind, + }, + }; + this.#runner.__sendToServer(message); + } + + closeActiveRequests(actor: RunnerActor) { + const actorId = actor.actorId; + + // Terminate all requests for this actor. This will no send a + // ToServerResponse* message since the actor will no longer be loaded. + // The gateway is responsible for closing the request. + for (const entry of actor.pendingRequests) { + entry.request.reject(new Error(`Actor ${actorId} stopped`)); + if (entry.gatewayId && entry.requestId) { + this.#removeRequestToActor(entry.gatewayId, entry.requestId); + } + } + + // Close all WebSockets. Only send close event to non-HWS. The gateway is + // responsible for hibernating HWS and closing regular WS. + for (const entry of actor.webSockets) { + const isHibernatable = entry.ws[HIBERNATABLE_SYMBOL]; + if (!isHibernatable) { + entry.ws._closeWithoutCallback(1000, "actor.stopped"); + } + // Note: request-to-actor mapping is cleaned up in the close callback + } + } + + async #fetch( + actorId: string, + gatewayId: protocol.GatewayId, + requestId: protocol.RequestId, + request: Request, + ): Promise { + // Validate actor exists + if (!this.#runner.hasActor(actorId)) { + this.log?.warn({ + msg: "ignoring request for unknown actor", + actorId, + }); + + // NOTE: This is a special response that will cause Guard to retry the request + // + // See should_retry_request_inner + // https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/guard-core/src/proxy_service.rs#L2458 + return new Response("Actor not found", { + status: 503, + headers: { "x-rivet-error": "runner.actor_not_found" }, + }); + } + + const fetchHandler = this.#runner.config.fetch( + this.#runner, + actorId, + gatewayId, + requestId, + request, + ); + + if (!fetchHandler) { + return new Response("Not Implemented", { status: 501 }); + } + + return fetchHandler; + } + + async handleTunnelMessage(message: protocol.ToClientTunnelMessage) { + // Parse the gateway ID, request ID, and message index from the messageId + const { gatewayId, requestId, messageIndex } = message.messageId; + + const gatewayIdStr = idToStr(gatewayId); + const requestIdStr = idToStr(requestId); + this.log?.debug({ + msg: "receive tunnel msg", + gatewayId: gatewayIdStr, + requestId: requestIdStr, + messageIndex: message.messageId.messageIndex, + message: stringifyToClientTunnelMessageKind(message.messageKind), + }); + + switch (message.messageKind.tag) { + case "ToClientRequestStart": + await this.#handleRequestStart( + gatewayId, + requestId, + message.messageKind.val, + ); + break; + case "ToClientRequestChunk": + await this.#handleRequestChunk( + gatewayId, + requestId, + message.messageKind.val, + ); + break; + case "ToClientRequestAbort": + await this.#handleRequestAbort(gatewayId, requestId); + break; + case "ToClientWebSocketOpen": + await this.#handleWebSocketOpen( + gatewayId, + requestId, + message.messageKind.val, + ); + break; + case "ToClientWebSocketMessage": { + await this.#handleWebSocketMessage( + gatewayId, + requestId, + messageIndex, + message.messageKind.val, + ); + break; + } + case "ToClientWebSocketClose": + await this.#handleWebSocketClose( + gatewayId, + requestId, + message.messageKind.val, + ); + break; + default: + unreachable(message.messageKind); + } + } + + async #handleRequestStart( + gatewayId: GatewayId, + requestId: RequestId, + req: protocol.ToClientRequestStart, + ) { + // Track this request for the actor + const requestIdStr = idToStr(requestId); + const actor = await this.#runner.getAndWaitForActor(req.actorId); + if (!actor) { + this.log?.warn({ + msg: "actor does not exist in handleRequestStart, request will leak", + actorId: req.actorId, + requestId: requestIdStr, + }); + return; + } + + // Add to request-to-actor mapping + this.addRequestToActor(gatewayId, requestId, req.actorId); + + try { + // Convert headers map to Headers object + const headers = new Headers(); + for (const [key, value] of req.headers) { + headers.append(key, value); + } + + // Create Request object + const request = new Request(`http://localhost${req.path}`, { + method: req.method, + headers, + body: req.body ? new Uint8Array(req.body) : undefined, + }); + + // Handle streaming request + if (req.stream) { + // Create a stream for the request body + const stream = new ReadableStream({ + start: (controller) => { + // Store controller for chunks + const existing = actor.getPendingRequest( + gatewayId, + requestId, + ); + if (existing) { + existing.streamController = controller; + existing.actorId = req.actorId; + existing.gatewayId = gatewayId; + existing.requestId = requestId; + } else { + actor.createPendingRequestWithStreamController( + gatewayId, + requestId, + 0, + controller, + ); + } + }, + }); + + // Create request with streaming body + const streamingRequest = new Request(request, { + body: stream, + duplex: "half", + } as any); + + // Call fetch handler with validation + const response = await this.#fetch( + req.actorId, + gatewayId, + requestId, + streamingRequest, + ); + await this.#sendResponse( + actor.actorId, + actor.generation, + gatewayId, + requestId, + response, + ); + } else { + // Non-streaming request + // Create a pending request entry to track messageIndex for the response + actor.createPendingRequest(gatewayId, requestId, 0); + + const response = await this.#fetch( + req.actorId, + gatewayId, + requestId, + request, + ); + await this.#sendResponse( + actor.actorId, + actor.generation, + gatewayId, + requestId, + response, + ); + } + } catch (error) { + if (error instanceof RunnerShutdownError) { + this.log?.debug({ msg: "catught runner shutdown error" }); + } else { + this.log?.error({ msg: "error handling request", error }); + this.#sendResponseError( + actor.actorId, + actor.generation, + gatewayId, + requestId, + 500, + "Internal Server Error", + ); + } + } finally { + // Clean up request tracking + if (this.#runner.hasActor(req.actorId, actor.generation)) { + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); + } + } + } + + async #handleRequestChunk( + gatewayId: GatewayId, + requestId: RequestId, + chunk: protocol.ToClientRequestChunk, + ) { + const actor = await this.getAndWaitForRequestActor( + gatewayId, + requestId, + ); + if (actor) { + const pending = actor.getPendingRequest(gatewayId, requestId); + if (pending?.streamController) { + pending.streamController.enqueue(new Uint8Array(chunk.body)); + if (chunk.finish) { + pending.streamController.close(); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); + } + } + } + } + + async #handleRequestAbort(gatewayId: GatewayId, requestId: RequestId) { + const actor = await this.getAndWaitForRequestActor( + gatewayId, + requestId, + ); + if (actor) { + const pending = actor.getPendingRequest(gatewayId, requestId); + if (pending?.streamController) { + pending.streamController.error(new Error("Request aborted")); + } + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); + } + } + + async #sendResponse( + actorId: string, + generation: number, + gatewayId: GatewayId, + requestId: ArrayBuffer, + response: Response, + ) { + if (!this.#runner.hasActor(actorId, generation)) { + this.log?.warn({ + msg: "actor not loaded to send response, assuming gateway has closed request", + actorId, + generation, + requestId, + }); + return; + } + + // Always treat responses as non-streaming for now + // In the future, we could detect streaming responses based on: + // - Transfer-Encoding: chunked + // - Content-Type: text/event-stream + // - Explicit stream flag from the handler + + // Read the body first to get the actual content + const body = response.body ? await response.arrayBuffer() : null; + + if (body && body.byteLength > MAX_PAYLOAD_SIZE) { + throw new Error("Response body too large"); + } + + // Convert headers to map and add Content-Length if not present + const headers = new Map(); + response.headers.forEach((value, key) => { + headers.set(key, value); + }); + + // Add Content-Length header if we have a body and it's not already set + if (body && !headers.has("content-length")) { + headers.set("content-length", String(body.byteLength)); + } + + // Send as non-streaming response if actor has not stopped + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerResponseStart", + val: { + status: response.status as protocol.u16, + headers, + body: body || null, + stream: false, + }, + }); + } + + #sendResponseError( + actorId: string, + generation: number, + gatewayId: GatewayId, + requestId: ArrayBuffer, + status: number, + message: string, + ) { + if (!this.#runner.hasActor(actorId, generation)) { + this.log?.warn({ + msg: "actor not loaded to send response, assuming gateway has closed request", + actorId, + generation, + requestId, + }); + return; + } + + const headers = new Map(); + headers.set("content-type", "text/plain"); + + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerResponseStart", + val: { + status: status as protocol.u16, + headers, + body: new TextEncoder().encode(message).buffer as ArrayBuffer, + stream: false, + }, + }); + } + + async #handleWebSocketOpen( + gatewayId: GatewayId, + requestId: RequestId, + open: protocol.ToClientWebSocketOpen, + ) { + // NOTE: This method is safe to be async since we will not receive any + // further WebSocket events until we send a ToServerWebSocketOpen + // tunnel message. We can do any async logic we need to between those two events. + // + // Sending a ToServerWebSocketClose will terminate the WebSocket early. + + const requestIdStr = idToStr(requestId); + + // Validate actor exists + const actor = await this.#runner.getAndWaitForActor(open.actorId); + if (!actor) { + this.log?.warn({ + msg: "ignoring websocket for unknown actor", + actorId: open.actorId, + }); + + // NOTE: Closing a WebSocket before open is equivalent to a Service + // Unavailable error and will cause Guard to retry the request + // + // See + // https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/pegboard-gateway/src/lib.rs#L238 + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketClose", + val: { + code: 1011, + reason: "Actor not found", + hibernate: false, + }, + }); + return; + } + + // Close existing WebSocket if one already exists for this request ID. + // This should never happen, but prevents any potential duplicate + // WebSockets from retransmits. + const existingAdapter = actor.getWebSocket(gatewayId, requestId); + if (existingAdapter) { + this.log?.warn({ + msg: "closing existing websocket for duplicate open event for the same request id", + requestId: requestIdStr, + }); + // Close without sending a message through the tunnel since the server + // already knows about the new connection + existingAdapter._closeWithoutCallback(1000, "ws.duplicate_open"); + } + + // Create WebSocket + try { + const request = buildRequestForWebSocket( + open.path, + Object.fromEntries(open.headers), + ); + + const canHibernate = + this.#runner.config.hibernatableWebSocket.canHibernate( + actor.actorId, + gatewayId, + requestId, + request, + ); + + // #createWebSocket will call `runner.config.websocket` under the + // hood to add the event listeners for open, etc. If this handler + // throws, then the WebSocket will be closed before sending the + // open event. + const adapter = await this.#createWebSocket( + actor.actorId, + gatewayId, + requestId, + requestIdStr, + 0, + canHibernate, + false, + request, + open.path, + Object.fromEntries(open.headers), + false, + ); + + // Create a PendingRequest entry to track the message index + actor.createPendingRequest(gatewayId, requestId, 0); + + // Open the WebSocket after `config.socket` so (a) the event + // handlers can be added and (b) any errors in `config.websocket` + // will cause the WebSocket to terminate before the open event. + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketOpen", + val: { + canHibernate, + }, + }); + + // Dispatch open event + adapter._handleOpen(requestId); + } catch (error) { + this.log?.error({ msg: "error handling websocket open", error }); + + // TODO: Call close event on adapter if needed + + // Send close on error + this.#sendMessage(gatewayId, requestId, { + tag: "ToServerWebSocketClose", + val: { + code: 1011, + reason: "Server Error", + hibernate: false, + }, + }); + + // Clean up actor tracking + actor.deleteWebSocket(gatewayId, requestId); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); + } + } + + async #handleWebSocketMessage( + gatewayId: GatewayId, + requestId: RequestId, + serverMessageIndex: number, + msg: protocol.ToClientWebSocketMessage, + ) { + const actor = await this.getAndWaitForRequestActor( + gatewayId, + requestId, + ); + if (actor) { + const adapter = actor.getWebSocket(gatewayId, requestId); + if (adapter) { + const data = msg.binary + ? new Uint8Array(msg.data) + : new TextDecoder().decode(new Uint8Array(msg.data)); + + adapter._handleMessage( + requestId, + data, + serverMessageIndex, + msg.binary, + ); + return; + } + } + + // TODO: This will never retransmit the socket and the socket will close + this.log?.warn({ + msg: "missing websocket for incoming websocket message, this may indicate the actor stopped before processing a message", + requestId, + }); + } + + sendHibernatableWebSocketMessageAck( + gatewayId: ArrayBuffer, + requestId: ArrayBuffer, + clientMessageIndex: number, + ) { + const requestIdStr = idToStr(requestId); + + this.log?.debug({ + msg: "ack ws msg", + requestId: requestIdStr, + index: clientMessageIndex, + }); + + if (clientMessageIndex < 0 || clientMessageIndex > 65535) + throw new Error("Invalid websocket ack index"); + + // Get the actor to find the gatewayId + // + // We don't have to wait for the actor to start since we're not calling + // any callbacks on the actor + const actor = this.getRequestActor(gatewayId, requestId); + if (!actor) { + this.log?.warn({ + msg: "cannot send websocket ack, actor not found", + requestId: requestIdStr, + }); + return; + } + + // Get gatewayId from the pending request + const pending = actor.getPendingRequest(gatewayId, requestId); + if (!pending?.gatewayId) { + this.log?.warn({ + msg: "cannot send websocket ack, gatewayId not found in pending request", + requestId: requestIdStr, + }); + return; + } + + // Send the ack message + this.#sendMessage(pending.gatewayId, requestId, { + tag: "ToServerWebSocketMessageAck", + val: { + index: clientMessageIndex, + }, + }); + } + + async #handleWebSocketClose( + gatewayId: GatewayId, + requestId: RequestId, + close: protocol.ToClientWebSocketClose, + ) { + const actor = await this.getAndWaitForRequestActor( + gatewayId, + requestId, + ); + if (actor) { + const adapter = actor.getWebSocket(gatewayId, requestId); + if (adapter) { + // We don't need to send a close response + adapter._handleClose( + requestId, + close.code || undefined, + close.reason || undefined, + ); + actor.deleteWebSocket(gatewayId, requestId); + actor.deletePendingRequest(gatewayId, requestId); + this.#removeRequestToActor(gatewayId, requestId); + } + } + } +} + +/** + * Builds a request that represents the incoming request for a given WebSocket. + * + * This request is not a real request and will never be sent. It's used to be passed to the actor to behave like a real incoming request. + */ +function buildRequestForWebSocket( + path: string, + headers: Record, +): Request { + // We need to manually ensure the original Upgrade/Connection WS + // headers are present + const fullHeaders = { + ...headers, + Upgrade: "websocket", + Connection: "Upgrade", + }; + + if (!path.startsWith("/")) { + throw new Error("Path must start with leading slash"); + } + + const request = new Request(`http://actor${path}`, { + method: "GET", + headers: fullHeaders, + }); + + return request; +} diff --git a/rivetkit-typescript/packages/engine-runner/src/utils.ts b/rivetkit-typescript/packages/engine-runner/src/utils.ts new file mode 100644 index 0000000000..a6ba6961b4 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/utils.ts @@ -0,0 +1,178 @@ +import { logger } from "./log"; + +// 20MiB. Keep in sync with runner_max_response_payload_body_size from engine/packages/config/src/config/pegboard.rs +export const MAX_PAYLOAD_SIZE = 20 * 1024 * 1024; + +export function unreachable(x: never): never { + throw `Unreachable: ${x}`; +} + +export interface BackoffOptions { + initialDelay?: number; + maxDelay?: number; + multiplier?: number; + jitter?: boolean; +} + +export function calculateBackoff( + attempt: number, + options: BackoffOptions = {}, +): number { + const { + initialDelay = 1000, + maxDelay = 30000, + multiplier = 2, + jitter = true, + } = options; + + let delay = Math.min(initialDelay * multiplier ** attempt, maxDelay); + + if (jitter) { + // Add random jitter between 0% and 25% of the delay + delay = delay * (1 + Math.random() * 0.25); + } + + return Math.floor(delay); +} + +export interface ParsedCloseReason { + group: string; + error: string; + rayId?: string; +} + +/** + * Parses a WebSocket close reason in the format: {group}.{error} or {group}.{error}#{ray_id} + * + * Examples: + * - "ws.eviction#t1s80so6h3irenp8ymzltfoittcl00" + * - "ws.client_closed" + * + * Returns undefined if the format is invalid + */ +export function parseWebSocketCloseReason( + reason: string, +): ParsedCloseReason | undefined { + const [mainPart, rayId] = reason.split("#"); + const [group, error] = mainPart.split("."); + + if (!group || !error) { + logger()?.warn({ msg: "failed to parse close reason", reason }); + return undefined; + } + + return { + group, + error, + rayId, + }; +} + +const U16_MAX = 65535; + +/** + * Wrapping greater than comparison for u16 values. + * Based on shared_state.rs wrapping_gt implementation. + */ +export function wrappingGtU16(a: number, b: number): boolean { + return a !== b && wrappingSub(a, b, U16_MAX) < U16_MAX / 2; +} + +/** + * Wrapping less than comparison for u16 values. + * Based on shared_state.rs wrapping_lt implementation. + */ +export function wrappingLtU16(a: number, b: number): boolean { + return a !== b && wrappingSub(b, a, U16_MAX) < U16_MAX / 2; +} + +/** + * Wrapping greater than or equal comparison for u16 values. + */ +export function wrappingGteU16(a: number, b: number): boolean { + return a === b || wrappingGtU16(a, b); +} + +/** + * Wrapping less than or equal comparison for u16 values. + */ +export function wrappingLteU16(a: number, b: number): boolean { + return a === b || wrappingLtU16(a, b); +} + +/** + * Performs wrapping addition for u16 values. + */ +export function wrappingAddU16(a: number, b: number): number { + return (a + b) % (U16_MAX + 1); +} + +/** + * Performs wrapping subtraction for u16 values. + */ +export function wrappingSubU16(a: number, b: number): number { + return wrappingSub(a, b, U16_MAX); +} + +/** + * Performs wrapping subtraction for unsigned integers. + */ +function wrappingSub(a: number, b: number, max: number): number { + const result = a - b; + if (result < 0) { + return result + max + 1; + } + return result; +} + +export function arraysEqual(a: ArrayBuffer, b: ArrayBuffer): boolean { + const ua = new Uint8Array(a); + const ub = new Uint8Array(b); + if (ua.length !== ub.length) return false; + for (let i = 0; i < ua.length; i++) { + if (ua[i] !== ub[i]) return false; + } + return true; +} + +/** + * Polyfill for Promise.withResolvers(). + * + * This is specifically for Cloudflare Workers. Their implementation of Promise.withResolvers does not work correctly. + */ +export function promiseWithResolvers(): { + promise: Promise; + resolve: (value: T | PromiseLike) => void; + reject: (reason?: any) => void; +} { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: any) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +export function idToStr(id: ArrayBuffer): string { + const bytes = new Uint8Array(id); + return Array.from(bytes) + .map((byte) => byte.toString(16).padStart(2, "0")) + .join(""); +} + +export function stringifyError(error: unknown): string { + if (error instanceof Error) { + return `${error.name}: ${error.message}${error.stack ? `\n${error.stack}` : ""}`; + } else if (typeof error === "string") { + return error; + } else if (typeof error === "object" && error !== null) { + try { + return `${JSON.stringify(error)}`; + } catch { + return `[object ${error.constructor?.name || "Object"}]`; + } + } else { + return String(error); + } +} diff --git a/rivetkit-typescript/packages/engine-runner/src/websocket-tunnel-adapter.ts b/rivetkit-typescript/packages/engine-runner/src/websocket-tunnel-adapter.ts new file mode 100644 index 0000000000..a40f0830d9 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/websocket-tunnel-adapter.ts @@ -0,0 +1,210 @@ +import type { Logger } from "pino"; +import { VirtualWebSocket, type UniversalWebSocket, type RivetMessageEvent } from "@rivetkit/virtual-websocket"; +import type { Tunnel } from "./tunnel"; +import { MAX_PAYLOAD_SIZE, wrappingAddU16, wrappingLteU16, wrappingSubU16 } from "./utils"; + +export const HIBERNATABLE_SYMBOL = Symbol("hibernatable"); + +export class WebSocketTunnelAdapter { + #readyState: 0 | 1 | 2 | 3 = 0; + #binaryType: "nodebuffer" | "arraybuffer" | "blob" = "nodebuffer"; + #ws: VirtualWebSocket; + #tunnel: Tunnel; + #actorId: string; + #requestId: string; + #hibernatable: boolean; + #serverMessageIndex: number; + #sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void; + #closeCallback: (code?: number, reason?: string) => void; + + get [HIBERNATABLE_SYMBOL](): boolean { + return this.#hibernatable; + } + + get #log(): Logger | undefined { + return this.#tunnel.log; + } + + constructor( + tunnel: Tunnel, + actorId: string, + requestId: string, + serverMessageIndex: number, + hibernatable: boolean, + isRestoringHibernatable: boolean, + public readonly request: Request, + sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void, + closeCallback: (code?: number, reason?: string) => void, + ) { + this.#tunnel = tunnel; + this.#actorId = actorId; + this.#requestId = requestId; + this.#hibernatable = hibernatable; + this.#serverMessageIndex = serverMessageIndex; + this.#sendCallback = sendCallback; + this.#closeCallback = closeCallback; + + this.#ws = new VirtualWebSocket({ + getReadyState: () => this.#readyState, + onSend: (data) => this.#handleSend(data), + onClose: (code, reason) => this.#close(code, reason, true), + onTerminate: () => this.#terminate(), + }); + + if (isRestoringHibernatable) { + this.#log?.debug({ + msg: "setting WebSocket to OPEN state for restored connection", + actorId: this.#actorId, + requestId: this.#requestId, + }); + this.#readyState = 1; + } + } + + get websocket(): UniversalWebSocket { + return this.#ws; + } + + #handleSend(data: string | ArrayBufferLike | Blob | ArrayBufferView): void { + let isBinary = false; + let messageData: string | ArrayBuffer; + + if (typeof data === "string") { + const encoder = new TextEncoder(); + if (encoder.encode(data).byteLength > MAX_PAYLOAD_SIZE) { + throw new Error("WebSocket message too large"); + } + + messageData = data; + } else if (data instanceof ArrayBuffer) { + if (data.byteLength > MAX_PAYLOAD_SIZE) throw new Error("WebSocket message too large"); + + isBinary = true; + messageData = data; + } else if (ArrayBuffer.isView(data)) { + if (data.byteLength > MAX_PAYLOAD_SIZE) throw new Error("WebSocket message too large"); + + isBinary = true; + const view = data; + const buffer = view.buffer instanceof SharedArrayBuffer + ? new Uint8Array(view.buffer, view.byteOffset, view.byteLength).slice().buffer + : view.buffer.slice(view.byteOffset, view.byteOffset + view.byteLength); + messageData = buffer as ArrayBuffer; + } else { + throw new Error("Unsupported data type"); + } + + this.#sendCallback(messageData, isBinary); + } + + // Called by Tunnel when WebSocket is opened + _handleOpen(requestId: ArrayBuffer): void { + if (this.#readyState !== 0) return; + this.#readyState = 1; + this.#ws.dispatchEvent({ type: "open", rivetRequestId: requestId, target: this.#ws }); + } + + // Called by Tunnel when message is received + _handleMessage( + requestId: ArrayBuffer, + data: string | Uint8Array, + serverMessageIndex: number, + isBinary: boolean, + ): boolean { + if (this.#readyState !== 1) { + this.#log?.warn({ + msg: "WebSocket message ignored - not in OPEN state", + requestId: this.#requestId, + actorId: this.#actorId, + currentReadyState: this.#readyState, + }); + return true; + } + + // Validate message index for hibernatable websockets + if (this.#hibernatable) { + const previousIndex = this.#serverMessageIndex; + + if (wrappingLteU16(serverMessageIndex, previousIndex)) { + this.#log?.info({ + msg: "received duplicate hibernating websocket message", + requestId, + actorId: this.#actorId, + previousIndex, + receivedIndex: serverMessageIndex, + }); + return true; + } + + const expectedIndex = wrappingAddU16(previousIndex, 1); + if (serverMessageIndex !== expectedIndex) { + const closeReason = "ws.message_index_skip"; + this.#log?.warn({ + msg: "hibernatable websocket message index out of sequence, closing connection", + requestId, + actorId: this.#actorId, + previousIndex, + expectedIndex, + receivedIndex: serverMessageIndex, + closeReason, + gap: wrappingSubU16(wrappingSubU16(serverMessageIndex, previousIndex), 1), + }); + this.#close(1008, closeReason, true); + return true; + } + + this.#serverMessageIndex = serverMessageIndex; + } + + // Convert data based on binaryType + let messageData: any = data; + if (isBinary && data instanceof Uint8Array) { + if (this.#binaryType === "nodebuffer") { + messageData = Buffer.from(data); + } else if (this.#binaryType === "arraybuffer") { + messageData = data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength); + } + } + + this.#ws.dispatchEvent({ + type: "message", + data: messageData, + rivetRequestId: requestId, + rivetMessageIndex: serverMessageIndex, + target: this.#ws, + } as RivetMessageEvent); + + return false; + } + + // Called by Tunnel when close is received + _handleClose(_requestId: ArrayBuffer, code?: number, reason?: string): void { + this.#close(code, reason, true); + } + + // Close without sending close message to tunnel + _closeWithoutCallback(code?: number, reason?: string): void { + this.#close(code, reason, false); + } + + // Public close method (used by tunnel.ts for stale websocket cleanup) + close(code?: number, reason?: string): void { + this.#close(code, reason, true); + } + + #close(code: number | undefined, reason: string | undefined, sendCallback: boolean): void { + if (this.#readyState >= 2) return; + + this.#readyState = 2; + if (sendCallback) this.#closeCallback(code, reason); + this.#readyState = 3; + this.#ws.triggerClose(code ?? 1000, reason ?? ""); + } + + #terminate(): void { + // Immediate close without close frame + this.#readyState = 3; + this.#closeCallback(1006, "Abnormal Closure"); + this.#ws.triggerClose(1006, "Abnormal Closure", false); + } +} diff --git a/rivetkit-typescript/packages/engine-runner/src/websocket.ts b/rivetkit-typescript/packages/engine-runner/src/websocket.ts new file mode 100644 index 0000000000..2804c77f50 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/src/websocket.ts @@ -0,0 +1,43 @@ +import { logger } from "./log"; + +// Global singleton promise that will be reused for subsequent calls +let webSocketPromise: Promise | null = null; + +export async function importWebSocket(): Promise { + // Return existing promise if we already started loading + if (webSocketPromise !== null) { + return webSocketPromise; + } + + // Create and store the promise + webSocketPromise = (async () => { + let _WebSocket: typeof WebSocket; + + if (typeof WebSocket !== "undefined") { + // Native + _WebSocket = WebSocket as unknown as typeof WebSocket; + logger()?.debug({ msg: "using native websocket" }); + } else { + // Node.js package + try { + const ws = await import("ws"); + _WebSocket = ws.default as unknown as typeof WebSocket; + logger()?.debug({ msg: "using websocket from npm" }); + } catch { + // WS not available + _WebSocket = class MockWebSocket { + constructor() { + throw new Error( + 'WebSocket support requires installing the "ws" peer dependency.', + ); + } + } as unknown as typeof WebSocket; + logger()?.debug({ msg: "using mock websocket" }); + } + } + + return _WebSocket; + })(); + + return webSocketPromise; +} diff --git a/rivetkit-typescript/packages/engine-runner/tests/lifecycle.test.ts b/rivetkit-typescript/packages/engine-runner/tests/lifecycle.test.ts new file mode 100644 index 0000000000..c84976890b --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/tests/lifecycle.test.ts @@ -0,0 +1,596 @@ +// import { describe, it, expect, vi } from "vitest"; +// import { Runner } from "@/mod"; +// import type { RunnerConfig, ActorConfig } from "@/mod"; +// import WebSocket, { type CloseEvent } from "ws"; +// +// const RIVET_ENDPOINT = process.env.RIVET_ENDPOINT ?? "http://localhost:6420"; +// const RIVET_ENDPOINT_WS = RIVET_ENDPOINT.replace("http://", "ws://").replace( +// "https://", +// "wss://", +// ); +// +// async function createActor( +// namespaceName: string, +// runnerNameSelector: string, +// durable: boolean, +// name?: string, +// ): Promise { +// const response = await fetch( +// `${RIVET_ENDPOINT}/actors?namespace=${namespaceName}`, +// { +// method: "POST", +// headers: { +// "Content-Type": "application/json", +// }, +// body: JSON.stringify({ +// name: name ?? "thingy", +// input: btoa("hello"), +// runner_name_selector: runnerNameSelector, +// durable, +// }), +// }, +// ); +// +// if (!response.ok) { +// throw new Error( +// `Failed to create actor: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// +// return response.json(); +// } +// +// async function destroyActor( +// namespaceName: string, +// actorId: string, +// ): Promise { +// const response = await fetch( +// `${RIVET_ENDPOINT}/actors/${actorId}?namespace=${namespaceName}`, +// { +// method: "DELETE", +// }, +// ); +// +// if (!response.ok) { +// throw new Error( +// `Failed to delete actor: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// } +// +// async function createNamespace( +// name: string, +// displayName: string, +// ): Promise { +// const response = await fetch(`${RIVET_ENDPOINT}/namespaces`, { +// method: "POST", +// headers: { +// "Content-Type": "application/json", +// }, +// body: JSON.stringify({ +// name, +// display_name: displayName, +// }), +// }); +// +// if (!response.ok) { +// console.warn( +// `Failed to create namespace: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// } +// +// async function getActorNames(namespaceName: string): Promise { +// const response = await fetch( +// `${RIVET_ENDPOINT}/actors/names?namespace=${namespaceName}`, +// { +// method: "GET", +// headers: { +// "Content-Type": "application/json", +// }, +// }, +// ); +// +// if (!response.ok) { +// throw new Error( +// `Failed to get actor names: ${response.status} ${response.statusText}\n${await response.text()}`, +// ); +// } +// +// return await response.json(); +// } +// +// describe("Runner E2E", () => { +// it("performs end-to-end actor lifecycle", async () => { +// const namespaceName = `test-${Math.floor(Math.random() * 10000)}`; +// const runnerName = "test-runner"; +// const prepopulateActorNames: string[] = Array.from( +// { length: 8 }, +// () => +// `actor-${Math.random().toString(36).substring(2, 10)}-${Date.now()}`, +// ); +// let runnerStarted = Promise.withResolvers(); +// let websocketOpen = Promise.withResolvers(); +// let websocketClosed = Promise.withResolvers(); +// let runner: Runner | null = null; +// const actorWebSockets = new Map(); +// +// // Use objects to hold the current promise resolvers so callbacks always get the latest +// const startedRef = { current: Promise.withResolvers() }; +// const stoppedRef = { current: Promise.withResolvers() }; +// +// const config: RunnerConfig = { +// version: 1, +// endpoint: RIVET_ENDPOINT, +// namespace: namespaceName, +// addresses: { main: { host: "127.0.0.1", port: 5051 } }, +// totalSlots: 100, +// runnerName: runnerName, +// runnerKey: "default", +// prepopulateActorNames, +// onConnected: () => { +// runnerStarted.resolve(undefined); +// }, +// onDisconnected: () => { }, +// fetch: async (actorId: string, request: Request) => { +// const url = new URL(request.url); +// if (url.pathname === "/ping") { +// // Return the actor ID in response +// return new Response( +// JSON.stringify({ +// actorId, +// status: "ok", +// timestamp: Date.now(), +// }), +// { +// status: 200, +// headers: { "Content-Type": "application/json" }, +// }, +// ); +// } +// return new Response("ok", { status: 200 }); +// }, +// onActorStart: async ( +// _actorId: string, +// _generation: number, +// _config: ActorConfig, +// ) => { +// console.log( +// `Actor ${_actorId} started (generation ${_generation})`, +// ); +// startedRef.current.resolve(undefined); +// }, +// onActorStop: async (_actorId: string, _generation: number) => { +// console.log( +// `Actor ${_actorId} stopped (generation ${_generation})`, +// ); +// stoppedRef.current.resolve(undefined); +// }, +// websocket: async ( +// actorId: string, +// ws: WebSocket, +// request: Request, +// ) => { +// console.log(`WebSocket connected for actor ${actorId}`); +// websocketOpen.resolve(undefined); +// actorWebSockets.set(actorId, ws); +// +// // Echo server - send back any messages received +// ws.on("message", (data) => { +// console.log( +// `WebSocket message from actor ${actorId}:`, +// data.toString(), +// ); +// ws.send(`Echo: ${data}`); +// }); +// +// ws.on("close", () => { +// console.log(`WebSocket closed for actor ${actorId}`); +// actorWebSockets.delete(actorId); +// websocketClosed.resolve(undefined); +// }); +// }, +// }; +// +// // Create namespace first +// await createNamespace(namespaceName, "Test Namespace"); +// +// runner = new Runner(config); +// +// // Check pegboard URL configuration +// expect(runner.pegboardUrl).toBe( +// `${RIVET_ENDPOINT_WS}/v1?namespace=${namespaceName}`, +// ); +// +// // Start runner +// runner.start(); +// +// // Wait for runner to be ready +// console.log("Waiting runner start..."); +// await runnerStarted.promise; +// +// // Check actor names prepopulated +// console.log("Comparing actor names..."); +// await vi.waitFor( +// async () => { +// const { names } = await getActorNames(namespaceName); +// expect(names.sort()).toStrictEqual( +// prepopulateActorNames.sort(), +// ); +// }, +// { interval: 100 }, +// ); +// +// // Create an actor +// console.log("Creating actor..."); +// const actorResponse = await createActor( +// namespaceName, +// runnerName, +// false, +// ); +// console.log("Actor created:", actorResponse.actor); +// const actorId = actorResponse.actor.actor_id; +// +// // Wait for actor to start +// console.log("Waiting new actor start..."); +// await startedRef.current.promise; +// +// // Ping actor to get actor ID in response (via Guard port) +// console.log("Pinging actor..."); +// const actorPingResponse = await fetch(`${RIVET_ENDPOINT}/ping`, { +// method: "GET", +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// expect(actorPingResponse.ok).toBe(true); +// const pingResult = (await actorPingResponse.json()) as any; +// expect(pingResult.actorId).toBe(actorId); +// +// // Test WebSocket connection +// console.log("Testing WebSocket connection..."); +// const ws = new WebSocket(`${RIVET_ENDPOINT_WS}/ws`, { +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// +// const testMessage = "Hello, actor!"; +// const messagePromise = new Promise((resolve, reject) => { +// ws.once("open", () => { +// console.log("WebSocket connected"); +// ws.send(testMessage); +// }); +// ws.once("message", (data) => { +// resolve(data.toString()); +// }); +// ws.once("error", reject); +// }); +// +// await websocketOpen.promise; +// +// // Test WebSocket messaging +// console.log("Testing WebSocket messaging..."); +// const response = await messagePromise; +// expect(response).toBe(`Echo: ${testMessage}`); +// +// // Close WebSocket for now +// ws.close(); +// console.log("Waiting websocket close..."); +// await websocketClosed.promise; +// +// await testKv(runner, actorId); +// +// // Sleep and wake actor 3 times in a loop +// for (let i = 1; i <= 3; i++) { +// console.log(`Sleep/wake cycle ${i}/3`); +// +// // Sleep actor +// console.log(`Sleeping actor (cycle ${i})...`); +// stoppedRef.current = Promise.withResolvers(); +// runner.sleepActor(actorId); +// +// console.log("Waiting actor sleep..."); +// await stoppedRef.current.promise; +// +// // Make network request to wake actor (via Guard) +// console.log(`Waking actor (cycle ${i})...`); +// startedRef.current = Promise.withResolvers(); +// const wakeResponse = await fetch(`${RIVET_ENDPOINT}/wake`, { +// method: "GET", +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// console.log(`Wake response status: ${wakeResponse.status}`); +// console.log(`Wake response body: ${await wakeResponse.text()}`); +// expect(wakeResponse.status).toBe(200); +// +// // TODO: Remove this +// // Wait for actor to wake +// console.log("Waiting actor start..."); +// await startedRef.current.promise; +// console.log(`Actor started successfully for cycle ${i}`); +// +// await testKvAfterSleep(runner, actorId); +// } +// +// // Sleep and wake actor 3 times in a loop +// for (let i = 1; i <= 3; i++) { +// console.log(`Sleep/wake cycle ${i}/3`); +// +// // Sleep actor +// console.log(`Sleeping actor (cycle ${i})...`); +// stoppedRef.current = Promise.withResolvers(); +// runner.sleepActor(actorId); +// +// console.log("Waiting actor sleep..."); +// await stoppedRef.current.promise; +// +// // Open websocket to wake actor (via Guard) +// console.log(`Waking actor (cycle ${i})...`); +// startedRef.current = Promise.withResolvers(); +// const ws = new WebSocket(`${RIVET_ENDPOINT_WS}/ws`, { +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// +// await new Promise((resolve, reject) => { +// ws.on("open", () => { +// console.log("WebSocket connected for wake test"); +// resolve(); +// }); +// ws.on("error", reject); +// }); +// +// // TODO: Remove this +// // Wait for actor to wake +// console.log("Waiting actor start..."); +// await startedRef.current.promise; +// console.log(`Actor started successfully for cycle ${i}`); +// +// await testKvAfterSleep(runner, actorId); +// } +// +// // Create a fresh WebSocket connection for destroy testing +// console.log("Creating WebSocket for destroy test..."); +// const wsForDestroy = new WebSocket(`${RIVET_ENDPOINT_WS}/ws`, { +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// +// await new Promise((resolve, reject) => { +// wsForDestroy.on("open", () => { +// console.log("WebSocket connected for destroy test"); +// resolve(); +// }); +// wsForDestroy.on("error", reject); +// }); +// +// // Test WebSocket closes on actor destroy +// const wsClosePromise = new Promise((resolve) => { +// wsForDestroy.on("close", () => { +// console.log("WebSocket closed after actor destroy"); +// resolve(); +// }); +// }); +// +// // Destroy actor +// console.log("Destroying actor..."); +// stoppedRef.current = Promise.withResolvers(); // Create new promise for actor destroy +// +// // Start destroy and wait for WebSocket close simultaneously +// const destroyPromise = destroyActor(namespaceName, actorId); +// +// // Wait for WebSocket to close +// console.log("Waiting WS close..."); +// await wsClosePromise; +// +// // Ensure destroy API call completed +// await destroyPromise; +// console.log("Destroy API call completed"); +// +// // Wait for actor to stop with timeout +// console.log("Waiting actor stopped..."); +// await stoppedRef.current.promise; +// console.log("Actor stop callback completed"); +// +// // Validate actor is destroyed +// console.log("Validating actor is destroyed..."); +// const destroyedPingResponse = await fetch(`${RIVET_ENDPOINT}/ping`, { +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// expect(destroyedPingResponse.status).toBe(404); +// +// // Test WebSocket connection to destroyed actor fails +// console.log("Testing WebSocket to destroyed actor..."); +// const wsToDestroyed = new WebSocket(`${RIVET_ENDPOINT_WS}/ws`, { +// headers: { +// "x-rivet-target": "actor", +// "x-rivet-actor": actorId, +// }, +// }); +// +// console.log( +// "Waiting WS close...", +// ); +// const closeCode = await new Promise((resolve, reject) => { +// wsToDestroyed.on("error", (err) => { +// console.log("WebSocket should not have errored"); +// reject(err); +// }); +// wsToDestroyed.on("close", (code) => { +// console.log("WebSocket closed"); +// resolve(code); +// }); +// }); +// expect(closeCode).toBe(1011); +// +// console.log("E2E test completed successfully!"); +// +// // Clean up - stop the runner +// if (runner) { +// await runner.shutdown(false); +// } +// }, 30_000); +// }); +// +// async function testKv(runner: Runner, actorId: string) { +// // Test KV operations +// console.log("Testing KV operations..."); +// +// // Test kvPut and kvGet +// const testEntries: [Uint8Array, Uint8Array][] = [ +// [createTestKey(["user", "123"]), createTestValue("alice")], +// [createTestKey(["user", "456"]), createTestValue("bob")], +// [createTestKey(["config", "theme"]), createTestValue("dark")], +// [createTestKey(["config", "lang"]), createTestValue("en")], +// ]; +// +// console.log("Testing kvPut..."); +// await runner.kvPut(actorId, testEntries); +// +// console.log("Testing kvGet..."); +// const getKeys = testEntries.map(([key, _]) => key); +// const getResult = await runner.kvGet(actorId, getKeys); +// +// expect(getResult.length).toBe(4); +// expect(decodeValue(getResult[0]!)).toBe("alice"); +// expect(decodeValue(getResult[1]!)).toBe("bob"); +// expect(decodeValue(getResult[2]!)).toBe("dark"); +// expect(decodeValue(getResult[3]!)).toBe("en"); +// +// // Test getting non-existent key +// const nonExistentResult = await runner.kvGet(actorId, [ +// createTestKey(["nonexistent"]), +// ]); +// expect(nonExistentResult[0]).toBe(null); +// +// // Test kvListAll +// console.log("Testing kvListAll..."); +// const allEntries = await runner.kvListAll(actorId); +// expect(allEntries.length).toBe(4); +// +// // Verify all entries are present (order may vary) +// const allValues = allEntries.map(([_, value]) => decodeValue(value)); +// expect(allValues.sort()).toEqual(["alice", "bob", "dark", "en"]); +// +// // Test kvListAll with limit +// const limitedEntries = await runner.kvListAll(actorId, { limit: 2 }); +// expect(limitedEntries.length).toBe(2); +// +// // Test kvListPrefix +// console.log("Testing kvListPrefix..."); +// const userEntries = await runner.kvListPrefix( +// actorId, +// createTestKey(["user"]), +// ); +// // Note: Prefix queries may not be working as expected on the server side +// // For now, we'll test that the method executes without error +// expect(userEntries.length).toBeGreaterThanOrEqual(0); +// +// const configEntries = await runner.kvListPrefix( +// actorId, +// createTestKey(["config"]), +// ); +// expect(configEntries.length).toBeGreaterThanOrEqual(0); +// +// // Test kvListRange +// console.log("Testing kvListRange..."); +// const rangeEntries = await runner.kvListRange( +// actorId, +// createTestKey(["config"]), +// createTestKey(["user"]), +// false, // inclusive +// ); +// // Range queries may have varying behavior depending on key ordering +// expect(rangeEntries.length).toBeGreaterThanOrEqual(0); +// +// // Test kvDelete +// console.log("Testing kvDelete..."); +// const keysToDelete = [createTestKey(["user", "456"])]; // Delete bob +// await runner.kvDelete(actorId, keysToDelete); +// +// // Verify deletion worked +// const afterDeleteResult = await runner.kvGet(actorId, [ +// createTestKey(["user", "456"]), +// ]); +// expect(afterDeleteResult[0]).toBe(null); +// +// // Verify other data still exists +// const remainingUserResult = await runner.kvGet(actorId, [ +// createTestKey(["user", "123"]), +// ]); +// expect(decodeValue(remainingUserResult[0]!)).toBe("alice"); +// +// // Test kvDrop operation before destroy +// console.log("Testing kvDrop..."); +// await runner.kvDrop(actorId); +// +// // Verify all data is cleared +// const afterDropData = await runner.kvGet(actorId, [ +// createTestKey(["user", "123"]), +// createTestKey(["config", "theme"]), +// createTestKey(["config", "lang"]), +// ]); +// expect(afterDropData[0]).toBe(null); +// expect(afterDropData[1]).toBe(null); +// expect(afterDropData[2]).toBe(null); +// +// // Verify list operations return empty after drop +// const afterDropList = await runner.kvListAll(actorId); +// expect(afterDropList.length).toBe(0); +// +// // Write data to test it exists during a sleep +// console.log("Writing data to live during sleep..."); +// await runner.kvPut(actorId, [ +// [createTestKey(["user", "789"]), createTestValue("max")], +// ]); +// } +// +// async function testKvAfterSleep(runner: Runner, actorId: string) { +// // Verify data still exists after waking again +// const remainingUserResult = await runner.kvGet(actorId, [ +// createTestKey(["user", "789"]), +// ]); +// expect(decodeValue(remainingUserResult[0]!)).toBe("max"); +// } +// +// function createTestKey(segments: string[]): Uint8Array { +// return flattenUint8Arrays(segments.map((s) => new TextEncoder().encode(s))); +// } +// +// function createTestValue(value: string): Uint8Array { +// return new TextEncoder().encode(value); +// } +// +// function decodeValue(value: Uint8Array): string { +// return new TextDecoder().decode(value); +// } +// +// function flattenUint8Arrays(arrays: Uint8Array[]): Uint8Array { +// // Calculate total length +// const totalLength = arrays.reduce((sum, arr) => sum + arr.length, 0); +// +// // Create result array +// const result = new Uint8Array(totalLength); +// +// // Copy each array +// let offset = 0; +// for (const arr of arrays) { +// result.set(arr, offset); +// offset += arr.length; +// } +// +// return result; +// } diff --git a/rivetkit-typescript/packages/engine-runner/tests/utils.test.ts b/rivetkit-typescript/packages/engine-runner/tests/utils.test.ts new file mode 100644 index 0000000000..6259921683 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/tests/utils.test.ts @@ -0,0 +1,194 @@ +import { describe, expect, it } from "vitest"; +import { + wrappingGteU16, + wrappingGtU16, + wrappingLteU16, + wrappingLtU16, +} from "../src/utils"; + +describe("wrappingGtU16", () => { + it("should return true when a > b in normal case", () => { + expect(wrappingGtU16(100, 50)).toBe(true); + expect(wrappingGtU16(1000, 999)).toBe(true); + }); + + it("should return false when a < b in normal case", () => { + expect(wrappingGtU16(50, 100)).toBe(false); + expect(wrappingGtU16(999, 1000)).toBe(false); + }); + + it("should return false when a == b", () => { + expect(wrappingGtU16(100, 100)).toBe(false); + expect(wrappingGtU16(0, 0)).toBe(false); + expect(wrappingGtU16(65535, 65535)).toBe(false); + }); + + it("should handle wrapping around u16 max", () => { + // When values wrap around, 1 is "greater than" 65535 + expect(wrappingGtU16(1, 65535)).toBe(true); + expect(wrappingGtU16(100, 65500)).toBe(true); + }); + + it("should handle edge cases near u16 boundaries", () => { + // 65535 is not greater than 0 (wrapped) + expect(wrappingGtU16(65535, 0)).toBe(false); + // But 0 is greater than 65535 if we consider wrapping + expect(wrappingGtU16(0, 65535)).toBe(true); + }); + + it("should handle values at exactly half the range", () => { + // U16_MAX / 2 = 32767.5, so values with distance <= 32767 return true + const lessThanHalf = 32766; + expect(wrappingGtU16(lessThanHalf, 0)).toBe(true); + expect(wrappingGtU16(0, lessThanHalf)).toBe(false); + + // At distance 32767, still less than 32767.5, so comparison returns true + const atHalfDistance = 32767; + expect(wrappingGtU16(atHalfDistance, 0)).toBe(true); + expect(wrappingGtU16(0, atHalfDistance)).toBe(false); + + // At distance 32768, greater than 32767.5, so comparison returns false + const overHalfDistance = 32768; + expect(wrappingGtU16(overHalfDistance, 0)).toBe(false); + expect(wrappingGtU16(0, overHalfDistance)).toBe(false); + }); +}); + +describe("wrappingLtU16", () => { + it("should return true when a < b in normal case", () => { + expect(wrappingLtU16(50, 100)).toBe(true); + expect(wrappingLtU16(999, 1000)).toBe(true); + }); + + it("should return false when a > b in normal case", () => { + expect(wrappingLtU16(100, 50)).toBe(false); + expect(wrappingLtU16(1000, 999)).toBe(false); + }); + + it("should return false when a == b", () => { + expect(wrappingLtU16(100, 100)).toBe(false); + expect(wrappingLtU16(0, 0)).toBe(false); + expect(wrappingLtU16(65535, 65535)).toBe(false); + }); + + it("should handle wrapping around u16 max", () => { + // When values wrap around, 65535 is "less than" 1 + expect(wrappingLtU16(65535, 1)).toBe(true); + expect(wrappingLtU16(65500, 100)).toBe(true); + }); + + it("should handle edge cases near u16 boundaries", () => { + // 0 is not less than 65535 (wrapped) + expect(wrappingLtU16(0, 65535)).toBe(false); + // But 65535 is less than 0 if we consider wrapping + expect(wrappingLtU16(65535, 0)).toBe(true); + }); + + it("should handle values at exactly half the range", () => { + // U16_MAX / 2 = 32767.5, so values with distance <= 32767 return true + const lessThanHalf = 32766; + expect(wrappingLtU16(0, lessThanHalf)).toBe(true); + expect(wrappingLtU16(lessThanHalf, 0)).toBe(false); + + // At distance 32767, still less than 32767.5, so comparison returns true + const atHalfDistance = 32767; + expect(wrappingLtU16(0, atHalfDistance)).toBe(true); + expect(wrappingLtU16(atHalfDistance, 0)).toBe(false); + + // At distance 32768, greater than 32767.5, so comparison returns false + const overHalfDistance = 32768; + expect(wrappingLtU16(0, overHalfDistance)).toBe(false); + expect(wrappingLtU16(overHalfDistance, 0)).toBe(false); + }); +}); + +describe("wrappingGtU16 and wrappingLtU16 consistency", () => { + it("should be inverse of each other for different values", () => { + const testCases: [number, number][] = [ + [100, 200], + [200, 100], + [0, 65535], + [65535, 0], + [1, 65534], + [32767, 32768], + ]; + + for (const [a, b] of testCases) { + const gt = wrappingGtU16(a, b); + const lt = wrappingLtU16(a, b); + const eq = a === b; + + // For any pair, exactly one of gt, lt, or eq should be true + expect(Number(gt) + Number(lt) + Number(eq)).toBe(1); + } + }); + + it("should satisfy transitivity for sequential values", () => { + // If we have sequential indices, a < b < c should hold + const a = 100; + const b = 101; + const c = 102; + + expect(wrappingLtU16(a, b)).toBe(true); + expect(wrappingLtU16(b, c)).toBe(true); + expect(wrappingLtU16(a, c)).toBe(true); + }); + + it("should handle sequence across wrap boundary", () => { + // Test a sequence that wraps: 65534, 65535, 0, 1 + const values = [65534, 65535, 0, 1]; + + for (let i = 0; i < values.length - 1; i++) { + expect(wrappingLtU16(values[i], values[i + 1])).toBe(true); + expect(wrappingGtU16(values[i + 1], values[i])).toBe(true); + } + }); +}); + +describe("wrappingGteU16", () => { + it("should return true when a > b", () => { + expect(wrappingGteU16(100, 50)).toBe(true); + expect(wrappingGteU16(1000, 999)).toBe(true); + }); + + it("should return true when a == b", () => { + expect(wrappingGteU16(100, 100)).toBe(true); + expect(wrappingGteU16(0, 0)).toBe(true); + expect(wrappingGteU16(65535, 65535)).toBe(true); + }); + + it("should return false when a < b", () => { + expect(wrappingGteU16(50, 100)).toBe(false); + expect(wrappingGteU16(999, 1000)).toBe(false); + }); + + it("should handle wrapping around u16 max", () => { + expect(wrappingGteU16(1, 65535)).toBe(true); + expect(wrappingGteU16(100, 65500)).toBe(true); + expect(wrappingGteU16(0, 65535)).toBe(true); + }); +}); + +describe("wrappingLteU16", () => { + it("should return true when a < b", () => { + expect(wrappingLteU16(50, 100)).toBe(true); + expect(wrappingLteU16(999, 1000)).toBe(true); + }); + + it("should return true when a == b", () => { + expect(wrappingLteU16(100, 100)).toBe(true); + expect(wrappingLteU16(0, 0)).toBe(true); + expect(wrappingLteU16(65535, 65535)).toBe(true); + }); + + it("should return false when a > b", () => { + expect(wrappingLteU16(100, 50)).toBe(false); + expect(wrappingLteU16(1000, 999)).toBe(false); + }); + + it("should handle wrapping around u16 max", () => { + expect(wrappingLteU16(65535, 1)).toBe(true); + expect(wrappingLteU16(65500, 100)).toBe(true); + expect(wrappingLteU16(65535, 0)).toBe(true); + }); +}); diff --git a/engine/sdks/typescript/test-envoy/tsconfig.json b/rivetkit-typescript/packages/engine-runner/tsconfig.json similarity index 65% rename from engine/sdks/typescript/test-envoy/tsconfig.json rename to rivetkit-typescript/packages/engine-runner/tsconfig.json index b74464ab70..959a71d91b 100644 --- a/engine/sdks/typescript/test-envoy/tsconfig.json +++ b/rivetkit-typescript/packages/engine-runner/tsconfig.json @@ -1,10 +1,10 @@ { - "extends": "../../../../tsconfig.base.json", + "extends": "../../../tsconfig.base.json", "compilerOptions": { + "types": ["node"], "paths": { "@/*": ["./src/*"] - }, - "lib": ["ESNext", "DOM"] + } }, "include": ["src/**/*", "tests/**/*", "benches/**/*"], "exclude": ["node_modules"] diff --git a/rivetkit-typescript/packages/engine-runner/tsup.config.ts b/rivetkit-typescript/packages/engine-runner/tsup.config.ts new file mode 100644 index 0000000000..f363b829fd --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/tsup.config.ts @@ -0,0 +1,4 @@ +import { defineConfig } from "tsup"; +import defaultConfig from "../../../tsup.base.ts"; + +export default defineConfig(defaultConfig); diff --git a/rivetkit-typescript/packages/engine-runner/turbo.json b/rivetkit-typescript/packages/engine-runner/turbo.json new file mode 100644 index 0000000000..29d4cb2625 --- /dev/null +++ b/rivetkit-typescript/packages/engine-runner/turbo.json @@ -0,0 +1,4 @@ +{ + "$schema": "https://turbo.build/schema.json", + "extends": ["//"] +} diff --git a/engine/sdks/typescript/test-envoy/vitest.config.ts b/rivetkit-typescript/packages/engine-runner/vitest.config.ts similarity index 83% rename from engine/sdks/typescript/test-envoy/vitest.config.ts rename to rivetkit-typescript/packages/engine-runner/vitest.config.ts index b6fc098098..146400235c 100644 --- a/engine/sdks/typescript/test-envoy/vitest.config.ts +++ b/rivetkit-typescript/packages/engine-runner/vitest.config.ts @@ -1,6 +1,6 @@ import { resolve } from "node:path"; import { defineConfig } from "vitest/config"; -import defaultConfig from "../../../../vitest.base.ts"; +import defaultConfig from "../../../vitest.base.ts"; export default defineConfig({ ...defaultConfig, diff --git a/rivetkit-typescript/packages/rivetkit/scripts/compile-bare.ts b/rivetkit-typescript/packages/rivetkit/scripts/compile-bare.ts index 965e0022cc..94d7ed9bf8 100755 --- a/rivetkit-typescript/packages/rivetkit/scripts/compile-bare.ts +++ b/rivetkit-typescript/packages/rivetkit/scripts/compile-bare.ts @@ -9,7 +9,7 @@ * 2. Replace Node.js assert import with a custom assert function * * IMPORTANT: Keep the post-processing logic in sync with: - * engine/sdks/rust/runner-protocol/build.rs + * engine/packages/runner-protocol/build.rs */ import * as fs from "node:fs/promises"; @@ -99,7 +99,7 @@ function assert(condition: boolean, message?: string): asserts condition { * 1. Replace @bare-ts/lib import with @rivetkit/bare-ts * 2. Replace Node.js assert import with a custom assert function * - * IMPORTANT: Keep this in sync with engine/sdks/rust/runner-protocol/build.rs + * IMPORTANT: Keep this in sync with engine/packages/runner-protocol/build.rs */ function postProcess(code: string): string { // Skip if already post-processed diff --git a/rivetkit-typescript/packages/sqlite-native/src/channel.rs b/rivetkit-typescript/packages/sqlite-native/src/channel.rs index aa62f0449d..de09d4cf2b 100644 --- a/rivetkit-typescript/packages/sqlite-native/src/channel.rs +++ b/rivetkit-typescript/packages/sqlite-native/src/channel.rs @@ -29,7 +29,7 @@ use crate::protocol::{ // MARK: Constants /// Timeout for individual KV operations in milliseconds. -/// Matches KV_EXPIRE in engine/sdks/typescript/runner/src/mod.ts. +/// Matches KV_EXPIRE in rivetkit-typescript/packages/engine-runner/src/mod.ts. const KV_EXPIRE_MS: u64 = 30_000; /// Initial reconnect delay in milliseconds. @@ -708,7 +708,7 @@ fn build_ws_url(config: &KvChannelConfig) -> String { /// Calculate exponential backoff delay with jitter. /// /// Matches the runner protocol reconnect strategy from -/// engine/sdks/typescript/runner/src/utils.ts. +/// rivetkit-typescript/packages/engine-runner/src/utils.ts. fn calculate_backoff(attempt: u32) -> Duration { let delay = INITIAL_BACKOFF_MS as f64 * BACKOFF_MULTIPLIER.powi(attempt as i32); let delay = delay.min(MAX_BACKOFF_MS as f64); @@ -881,7 +881,7 @@ mod tests { #[test] fn backoff_constants_match_runner_protocol() { - // These must match engine/sdks/typescript/runner/src/utils.ts. + // These must match rivetkit-typescript/packages/engine-runner/src/utils.ts. assert_eq!(INITIAL_BACKOFF_MS, 1000); assert_eq!(MAX_BACKOFF_MS, 30_000); assert!((BACKOFF_MULTIPLIER - 2.0).abs() < f64::EPSILON); diff --git a/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts b/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts index ebb6c86a92..5f510a8077 100644 --- a/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts +++ b/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts @@ -9,7 +9,7 @@ * 2. Replace Node.js assert import with a custom assert function * * IMPORTANT: Keep the post-processing logic in sync with: - * engine/sdks/rust/runner-protocol/build.rs + * engine/packages/runner-protocol/build.rs */ import * as fs from "node:fs/promises"; @@ -97,7 +97,7 @@ function assert(condition: boolean, message?: string): asserts condition { * 1. Replace @bare-ts/lib import with @rivetkit/bare-ts * 2. Replace Node.js assert import with a custom assert function * - * IMPORTANT: Keep this in sync with engine/sdks/rust/runner-protocol/build.rs + * IMPORTANT: Keep this in sync with engine/packages/runner-protocol/build.rs */ function postProcess(code: string): string { // Skip if already post-processed diff --git a/scripts/release/sdk.ts b/scripts/release/sdk.ts index 98893207c7..0b29cf8351 100644 --- a/scripts/release/sdk.ts +++ b/scripts/release/sdk.ts @@ -106,8 +106,8 @@ export async function publishSdk(opts: ReleaseOpts) { // Get list of packages to publish const enginePackagePaths = [ - `${opts.root}/engine/sdks/typescript/runner`, - `${opts.root}/engine/sdks/typescript/runner-protocol`, + `${opts.root}/rivetkit-typescript/packages/engine-runner`, + `${opts.root}/rivetkit-typescript/packages/engine-runner-protocol`, `${opts.root}/engine/sdks/typescript/api-full`, ]; diff --git a/scripts/tests/actor_e2e.ts b/scripts/tests/actor_e2e.ts index 75395a4aae..3000115e55 100755 --- a/scripts/tests/actor_e2e.ts +++ b/scripts/tests/actor_e2e.ts @@ -10,6 +10,7 @@ import { async function main() { let actorId; + const runnerNameSelector = getRunnerNameSelector(); try { console.log("Starting actor E2E test..."); @@ -22,7 +23,10 @@ async function main() { // Create an actor console.log("Creating actor..."); console.time("actor create"); - const actorResponse = await createActor(RIVET_NAMESPACE, "test"); + const actorResponse = await createActor( + RIVET_NAMESPACE, + runnerNameSelector, + ); console.timeEnd("actor create"); console.log("Actor created:", actorResponse.actor); @@ -77,6 +81,35 @@ async function main() { } } +function getRunnerNameSelector(): string { + if (process.env.RUNNER_NAME_SELECTOR) { + return process.env.RUNNER_NAME_SELECTOR; + } + + const implementation = ( + process.env.TEST_ENVOY_IMPL ?? + process.env.ACTOR_E2E_TEST_ENVOY_IMPL ?? + "rust" + ).toLowerCase(); + + if (implementation === "rust") { + return "test-envoy"; + } + + if ( + implementation === "typescript" || + implementation === "ts" || + implementation === "node" || + implementation === "napi" + ) { + return "test-envoy-ts"; + } + + throw new Error( + `Unsupported test envoy implementation: ${implementation}. Expected "rust" or "typescript".`, + ); +} + function testWebSocket(actorId: string): Promise { console.log("Testing WebSocket connection to actor...");