From 29cfbf98ebd32f0b50ebb1037855a94cb8859d11 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 3 Sep 2025 16:59:03 -0700 Subject: [PATCH] chore(runner): add ability to disable auto-shutdown --- sdks/typescript/runner/src/mod.ts | 33 +++++++---- sdks/typescript/test-runner/package.json | 1 + sdks/typescript/test-runner/src/main.ts | 71 ++++++++++++------------ 3 files changed, 58 insertions(+), 47 deletions(-) diff --git a/sdks/typescript/runner/src/mod.ts b/sdks/typescript/runner/src/mod.ts index 2ef6381721..5d35217478 100644 --- a/sdks/typescript/runner/src/mod.ts +++ b/sdks/typescript/runner/src/mod.ts @@ -34,17 +34,14 @@ export interface RunnerConfig { onConnected: () => void; onDisconnected: () => void; fetch: (actorId: string, request: Request) => Promise; - websocket?: ( - actorId: string, - ws: any, - request: Request, - ) => Promise; + websocket?: (actorId: string, ws: any, request: Request) => Promise; onActorStart: ( actorId: string, generation: number, config: ActorConfig, ) => Promise; onActorStop: (actorId: string, generation: number) => Promise; + noAutoShutdown?: boolean; } export interface KvListOptions { @@ -221,8 +218,10 @@ export class Runner { await this.#openPegboardWebSocket(); this.#openTunnel(); - process.on("SIGTERM", this.shutdown.bind(this, false, true)); - process.on("SIGINT", this.shutdown.bind(this, false, true)); + if (!this.#config.noAutoShutdown) { + process.on("SIGTERM", this.shutdown.bind(this, false, true)); + process.on("SIGINT", this.shutdown.bind(this, false, true)); + } } // MARK: Shutdown @@ -344,6 +343,10 @@ export class Runner { this.#tunnel.shutdown(); console.log("Tunnel shutdown completed"); } + + if (exit) { + process.exit(0); + } } // MARK: Networking @@ -356,7 +359,10 @@ export class Runner { } get pegboardRelayUrl() { - const endpoint = this.#config.pegboardRelayEndpoint || this.#config.pegboardEndpoint || this.#config.endpoint; + const endpoint = + this.#config.pegboardRelayEndpoint || + this.#config.pegboardEndpoint || + this.#config.endpoint; const wsEndpoint = endpoint .replace("http://", "ws://") .replace("https://", "wss://"); @@ -424,7 +430,7 @@ export class Runner { key: this.#config.runnerKey, version: this.#config.version, totalSlots: this.#config.totalSlots, - addressesHttp: new Map(), // No addresses needed with tunnel + addressesHttp: new Map(), // No addresses needed with tunnel addressesTcp: null, addressesUdp: null, lastCommandIdx: @@ -499,7 +505,9 @@ export class Runner { this.runnerId = init.runnerId; // Store the runner lost threshold from metadata - this.#runnerLostThreshold = init.metadata?.runnerLostThreshold ? Number(init.metadata.runnerLostThreshold) : undefined; + this.#runnerLostThreshold = init.metadata?.runnerLostThreshold + ? Number(init.metadata.runnerLostThreshold) + : undefined; console.log("Received init", { runnerId: init.runnerId, @@ -616,7 +624,10 @@ export class Runner { console.log("[RUNNER] Registering new actor with tunnel:", actorId); this.#tunnel.registerActor(actorId); } else { - console.error("[RUNNER] WARNING: No tunnel available to register actor:", actorId); + console.error( + "[RUNNER] WARNING: No tunnel available to register actor:", + actorId, + ); } this.#sendActorStateUpdate(actorId, generation, "running"); diff --git a/sdks/typescript/test-runner/package.json b/sdks/typescript/test-runner/package.json index 697d656dd2..7092ba4acf 100644 --- a/sdks/typescript/test-runner/package.json +++ b/sdks/typescript/test-runner/package.json @@ -3,6 +3,7 @@ "version": "0.0.0", "type": "module", "scripts": { + "start": "tsx src/main.ts", "build": "tsup src/main.ts", "check-types": "tsc --noEmit" }, diff --git a/sdks/typescript/test-runner/src/main.ts b/sdks/typescript/test-runner/src/main.ts index 9ca7f04ee7..596dda5009 100644 --- a/sdks/typescript/test-runner/src/main.ts +++ b/sdks/typescript/test-runner/src/main.ts @@ -3,11 +3,18 @@ import type { RunnerConfig, ActorConfig } from "@rivetkit/engine-runner"; import WebSocket from "ws"; import { serve } from "@hono/node-server"; -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_RUNNER_KEY = process.env.RIVET_RUNNER_KEY ?? `key-${Math.floor(Math.random() * 10000)}`; -const RIVET_RUNNER_VERSION = process.env.RIVET_RUNNER_VERSION ? Number(process.env.RIVET_RUNNER_VERSION) : 1; -const RIVET_RUNNER_TOTAL_SLOTS = process.env.RIVET_RUNNER_TOTAL_SLOTS ? Number(process.env.RIVET_RUNNER_TOTAL_SLOTS) : 100; +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_RUNNER_KEY = + process.env.RIVET_RUNNER_KEY ?? `key-${Math.floor(Math.random() * 10000)}`; +const RIVET_RUNNER_VERSION = process.env.RIVET_RUNNER_VERSION + ? Number(process.env.RIVET_RUNNER_VERSION) + : 1; +const RIVET_RUNNER_TOTAL_SLOTS = process.env.RIVET_RUNNER_TOTAL_SLOTS + ? Number(process.env.RIVET_RUNNER_TOTAL_SLOTS) + : 100; const RIVET_ENDPOINT = process.env.RIVET_ENDPOINT ?? "http://localhost:6420"; let runnerStarted = Promise.withResolvers(); @@ -20,18 +27,22 @@ const actorWebSockets = new Map(); serve({ fetch: async (request: Request) => { const url = new URL(request.url); - if (url.pathname == '/wait-ready') { + if (url.pathname == "/wait-ready") { await runnerStarted.promise; - return new Response(JSON.stringify(runner?.runnerId), { status: 200 }); - } else if (url.pathname == '/has-actor') { - let actorIdQuery = url.searchParams.get('actor'); - let generationQuery = url.searchParams.get('generation'); - let generation = generationQuery ? Number(generationQuery) : undefined; + return new Response(JSON.stringify(runner?.runnerId), { + status: 200, + }); + } else if (url.pathname == "/has-actor") { + let actorIdQuery = url.searchParams.get("actor"); + let generationQuery = url.searchParams.get("generation"); + let generation = generationQuery + ? Number(generationQuery) + : undefined; if (!actorIdQuery || !runner?.hasActor(actorIdQuery, generation)) { return new Response(undefined, { status: 404 }); } - } else if (url.pathname == '/shutdown') { + } else if (url.pathname == "/shutdown") { await runner?.shutdown(true); } @@ -56,9 +67,11 @@ const config: RunnerConfig = { onConnected: () => { runnerStarted.resolve(undefined); }, - onDisconnected: () => { }, + onDisconnected: () => {}, fetch: async (actorId: string, request: Request) => { - console.log(`[TEST-RUNNER] Fetch called for actor ${actorId}, URL: ${request.url}`); + console.log( + `[TEST-RUNNER] Fetch called for actor ${actorId}, URL: ${request.url}`, + ); const url = new URL(request.url); if (url.pathname === "/ping") { // Return the actor ID in response @@ -68,13 +81,10 @@ const config: RunnerConfig = { timestamp: Date.now(), }; console.log(`[TEST-RUNNER] Returning ping response:`, responseData); - return new Response( - JSON.stringify(responseData), - { - status: 200, - headers: { "Content-Type": "application/json" }, - }, - ); + return new Response(JSON.stringify(responseData), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); } return new Response("ok", { status: 200 }); @@ -84,22 +94,14 @@ const config: RunnerConfig = { _generation: number, _config: ActorConfig, ) => { - console.log( - `Actor ${_actorId} started (generation ${_generation})`, - ); + console.log(`Actor ${_actorId} started (generation ${_generation})`); startedRef.current.resolve(undefined); }, onActorStop: async (_actorId: string, _generation: number) => { - console.log( - `Actor ${_actorId} stopped (generation ${_generation})`, - ); + console.log(`Actor ${_actorId} stopped (generation ${_generation})`); stoppedRef.current.resolve(undefined); }, - websocket: async ( - actorId: string, - ws: WebSocket, - request: Request, - ) => { + websocket: async (actorId: string, ws: WebSocket, request: Request) => { console.log(`WebSocket connected for actor ${actorId}`); websocketOpen.resolve(undefined); actorWebSockets.set(actorId, ws); @@ -107,10 +109,7 @@ const config: RunnerConfig = { // Echo server - send back any messages received ws.addEventListener("message", (event) => { const data = event.data; - console.log( - `WebSocket message from actor ${actorId}:`, - data, - ); + console.log(`WebSocket message from actor ${actorId}:`, data); ws.send(`Echo: ${data}`); });