Skip to content

Commit 6ace4ef

Browse files
NathanFlurryNathanFlurry
committed
chore(runner): add ability to disable auto-shutdown
1 parent 6cf9da6 commit 6ace4ef

File tree

3 files changed

+58
-47
lines changed

3 files changed

+58
-47
lines changed

‎sdks/typescript/runner/src/mod.ts‎

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -34,17 +34,14 @@ export interface RunnerConfig {
3434
onConnected: () => void;
3535
onDisconnected: () => void;
3636
fetch: (actorId: string, request: Request) => Promise<Response>;
37-
websocket?: (
38-
actorId: string,
39-
ws: any,
40-
request: Request,
41-
) => Promise<void>;
37+
websocket?: (actorId: string, ws: any, request: Request) => Promise<void>;
4238
onActorStart: (
4339
actorId: string,
4440
generation: number,
4541
config: ActorConfig,
4642
) => Promise<void>;
4743
onActorStop: (actorId: string, generation: number) => Promise<void>;
44+
noAutoShutdown?: boolean;
4845
}
4946

5047
export interface KvListOptions {
@@ -221,8 +218,10 @@ export class Runner {
221218
await this.#openPegboardWebSocket();
222219
this.#openTunnel();
223220

224-
process.on("SIGTERM", this.shutdown.bind(this, false, true));
225-
process.on("SIGINT", this.shutdown.bind(this, false, true));
221+
if (!this.#config.noAutoShutdown) {
222+
process.on("SIGTERM", this.shutdown.bind(this, false, true));
223+
process.on("SIGINT", this.shutdown.bind(this, false, true));
224+
}
226225
}
227226

228227
// MARK: Shutdown
@@ -344,6 +343,10 @@ export class Runner {
344343
this.#tunnel.shutdown();
345344
console.log("Tunnel shutdown completed");
346345
}
346+
347+
if (exit) {
348+
process.exit(0);
349+
}
347350
}
348351

349352
// MARK: Networking
@@ -356,7 +359,10 @@ export class Runner {
356359
}
357360

358361
get pegboardRelayUrl() {
359-
const endpoint = this.#config.pegboardRelayEndpoint || this.#config.pegboardEndpoint || this.#config.endpoint;
362+
const endpoint =
363+
this.#config.pegboardRelayEndpoint ||
364+
this.#config.pegboardEndpoint ||
365+
this.#config.endpoint;
360366
const wsEndpoint = endpoint
361367
.replace("http://", "ws://")
362368
.replace("https://", "wss://");
@@ -424,7 +430,7 @@ export class Runner {
424430
key: this.#config.runnerKey,
425431
version: this.#config.version,
426432
totalSlots: this.#config.totalSlots,
427-
addressesHttp: new Map(), // No addresses needed with tunnel
433+
addressesHttp: new Map(), // No addresses needed with tunnel
428434
addressesTcp: null,
429435
addressesUdp: null,
430436
lastCommandIdx:
@@ -499,7 +505,9 @@ export class Runner {
499505
this.runnerId = init.runnerId;
500506

501507
// Store the runner lost threshold from metadata
502-
this.#runnerLostThreshold = init.metadata?.runnerLostThreshold ? Number(init.metadata.runnerLostThreshold) : undefined;
508+
this.#runnerLostThreshold = init.metadata?.runnerLostThreshold
509+
? Number(init.metadata.runnerLostThreshold)
510+
: undefined;
503511

504512
console.log("Received init", {
505513
runnerId: init.runnerId,
@@ -616,7 +624,10 @@ export class Runner {
616624
console.log("[RUNNER] Registering new actor with tunnel:", actorId);
617625
this.#tunnel.registerActor(actorId);
618626
} else {
619-
console.error("[RUNNER] WARNING: No tunnel available to register actor:", actorId);
627+
console.error(
628+
"[RUNNER] WARNING: No tunnel available to register actor:",
629+
actorId,
630+
);
620631
}
621632

622633
this.#sendActorStateUpdate(actorId, generation, "running");

‎sdks/typescript/test-runner/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"version": "0.0.0",
44
"type": "module",
55
"scripts": {
6+
"start": "tsx src/main.ts",
67
"build": "tsup src/main.ts",
78
"check-types": "tsc --noEmit"
89
},

‎sdks/typescript/test-runner/src/main.ts‎

Lines changed: 35 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,18 @@ import type { RunnerConfig, ActorConfig } from "@rivetkit/engine-runner";
33
import WebSocket from "ws";
44
import { serve } from "@hono/node-server";
55

6-
const INTERNAL_SERVER_PORT = process.env.INTERNAL_SERVER_PORT ? Number(process.env.INTERNAL_SERVER_PORT) : 5051;
7-
const RIVET_NAMESPACE = process.env.RIVET_NAMESPACE ?? 'default';
8-
const RIVET_RUNNER_KEY = process.env.RIVET_RUNNER_KEY ?? `key-${Math.floor(Math.random() * 10000)}`;
9-
const RIVET_RUNNER_VERSION = process.env.RIVET_RUNNER_VERSION ? Number(process.env.RIVET_RUNNER_VERSION) : 1;
10-
const RIVET_RUNNER_TOTAL_SLOTS = process.env.RIVET_RUNNER_TOTAL_SLOTS ? Number(process.env.RIVET_RUNNER_TOTAL_SLOTS) : 100;
6+
const INTERNAL_SERVER_PORT = process.env.INTERNAL_SERVER_PORT
7+
? Number(process.env.INTERNAL_SERVER_PORT)
8+
: 5051;
9+
const RIVET_NAMESPACE = process.env.RIVET_NAMESPACE ?? "default";
10+
const RIVET_RUNNER_KEY =
11+
process.env.RIVET_RUNNER_KEY ?? `key-${Math.floor(Math.random() * 10000)}`;
12+
const RIVET_RUNNER_VERSION = process.env.RIVET_RUNNER_VERSION
13+
? Number(process.env.RIVET_RUNNER_VERSION)
14+
: 1;
15+
const RIVET_RUNNER_TOTAL_SLOTS = process.env.RIVET_RUNNER_TOTAL_SLOTS
16+
? Number(process.env.RIVET_RUNNER_TOTAL_SLOTS)
17+
: 100;
1118
const RIVET_ENDPOINT = process.env.RIVET_ENDPOINT ?? "http://localhost:6420";
1219

1320
let runnerStarted = Promise.withResolvers();
@@ -20,18 +27,22 @@ const actorWebSockets = new Map<string, WebSocket>();
2027
serve({
2128
fetch: async (request: Request) => {
2229
const url = new URL(request.url);
23-
if (url.pathname == '/wait-ready') {
30+
if (url.pathname == "/wait-ready") {
2431
await runnerStarted.promise;
25-
return new Response(JSON.stringify(runner?.runnerId), { status: 200 });
26-
} else if (url.pathname == '/has-actor') {
27-
let actorIdQuery = url.searchParams.get('actor');
28-
let generationQuery = url.searchParams.get('generation');
29-
let generation = generationQuery ? Number(generationQuery) : undefined;
32+
return new Response(JSON.stringify(runner?.runnerId), {
33+
status: 200,
34+
});
35+
} else if (url.pathname == "/has-actor") {
36+
let actorIdQuery = url.searchParams.get("actor");
37+
let generationQuery = url.searchParams.get("generation");
38+
let generation = generationQuery
39+
? Number(generationQuery)
40+
: undefined;
3041

3142
if (!actorIdQuery || !runner?.hasActor(actorIdQuery, generation)) {
3243
return new Response(undefined, { status: 404 });
3344
}
34-
} else if (url.pathname == '/shutdown') {
45+
} else if (url.pathname == "/shutdown") {
3546
await runner?.shutdown(true);
3647
}
3748

@@ -56,9 +67,11 @@ const config: RunnerConfig = {
5667
onConnected: () => {
5768
runnerStarted.resolve(undefined);
5869
},
59-
onDisconnected: () => { },
70+
onDisconnected: () => {},
6071
fetch: async (actorId: string, request: Request) => {
61-
console.log(`[TEST-RUNNER] Fetch called for actor ${actorId}, URL: ${request.url}`);
72+
console.log(
73+
`[TEST-RUNNER] Fetch called for actor ${actorId}, URL: ${request.url}`,
74+
);
6275
const url = new URL(request.url);
6376
if (url.pathname === "/ping") {
6477
// Return the actor ID in response
@@ -68,13 +81,10 @@ const config: RunnerConfig = {
6881
timestamp: Date.now(),
6982
};
7083
console.log(`[TEST-RUNNER] Returning ping response:`, responseData);
71-
return new Response(
72-
JSON.stringify(responseData),
73-
{
74-
status: 200,
75-
headers: { "Content-Type": "application/json" },
76-
},
77-
);
84+
return new Response(JSON.stringify(responseData), {
85+
status: 200,
86+
headers: { "Content-Type": "application/json" },
87+
});
7888
}
7989

8090
return new Response("ok", { status: 200 });
@@ -84,33 +94,22 @@ const config: RunnerConfig = {
8494
_generation: number,
8595
_config: ActorConfig,
8696
) => {
87-
console.log(
88-
`Actor ${_actorId} started (generation ${_generation})`,
89-
);
97+
console.log(`Actor ${_actorId} started (generation ${_generation})`);
9098
startedRef.current.resolve(undefined);
9199
},
92100
onActorStop: async (_actorId: string, _generation: number) => {
93-
console.log(
94-
`Actor ${_actorId} stopped (generation ${_generation})`,
95-
);
101+
console.log(`Actor ${_actorId} stopped (generation ${_generation})`);
96102
stoppedRef.current.resolve(undefined);
97103
},
98-
websocket: async (
99-
actorId: string,
100-
ws: WebSocket,
101-
request: Request,
102-
) => {
104+
websocket: async (actorId: string, ws: WebSocket, request: Request) => {
103105
console.log(`WebSocket connected for actor ${actorId}`);
104106
websocketOpen.resolve(undefined);
105107
actorWebSockets.set(actorId, ws);
106108

107109
// Echo server - send back any messages received
108110
ws.addEventListener("message", (event) => {
109111
const data = event.data;
110-
console.log(
111-
`WebSocket message from actor ${actorId}:`,
112-
data,
113-
);
112+
console.log(`WebSocket message from actor ${actorId}:`, data);
114113
ws.send(`Echo: ${data}`);
115114
});
116115

0 commit comments

Comments
 (0)