fix: fix serverless config req, handle serverless sse close (#1368)

Author: MasterPtato

examples/smoke-test/scripts/connect.ts

@@ -15,8 +15,24 @@ async function main() {
 		await new Promise((resolve) => setTimeout(resolve, 1000));
 	}
 
-	await new Promise((resolve) => setTimeout(resolve, 10000));
+	await new Promise((resolve) => setTimeout(resolve, 2000));
 	await counter.dispose();
+
+	await new Promise((resolve) => setTimeout(resolve, 200));
+
+	const counter2 = client.counter.getOrCreate().connect();
+
+	counter2.on("newCount", (count: number) => console.log("Event:", count));
+
+	for (let i = 0; i < 5; i++) {
+		const out = await counter2.increment(5);
+		console.log("RPC:", out);
+
+		await new Promise((resolve) => setTimeout(resolve, 1000));
+	}
+
+	await new Promise((resolve) => setTimeout(resolve, 2000));
+	await counter2.dispose();
 }
 
 main();

examples/smoke-test/src/server/server.ts

@@ -1,3 +1,7 @@
 import { registry } from "./registry";
 
-registry.start();
+registry.start({
+	// defaultServerPort: 8080,
+	// runnerKind: "serverless",
+	// autoConfigureServerless: true,
+});

examples/smoke-test/src/smoke-test/spawn-actor.ts

@@ -37,13 +37,19 @@ export async function spawnActor({
 		await counter.increment(1);
 		await counter.dispose();
 
-		// Wait for actor to sleep
-		await new Promise(res => setTimeout(res, 1000));
-
+		// Immediately reconnect
 		const counter2 = client.counter.getOrCreate(key).connect();
 		await counter2.increment(1);
 		await counter2.dispose();
 
+		// Wait for actor to sleep
+		await new Promise(res => setTimeout(res, 1000));
+
+		// Reconnect after sleep
+		const counter3 = client.counter.getOrCreate(key).connect();
+		await counter3.increment(1);
+		await counter3.dispose();
+
 		succeeded = true;
 		onSuccess();
 	} catch (error) {

packages/rivetkit/src/drivers/engine/actor-driver.ts

@@ -411,6 +411,8 @@ export class EngineActorDriver implements ActorDriver {
 		await this.#runnerStarted.promise;
 
 		return streamSSE(c, async (stream) => {
+			stream.onAbort(() => this.shutdown(true));
+
 			// Runner id should be set if the runner started
 			const payload = this.#runner.getServerlessInitPacket();
 			invariant(payload, "runnerId not set");

packages/rivetkit/src/registry/mod.ts

@@ -92,6 +92,7 @@ export class Registry<A extends RegistryActors> {
 		if (config.runnerKind === "serverless") {
 			config.defaultServerPort = 8080;
 			config.overrideServerAddress = config.endpoint;
+			config.disableActorDriver = true;
 		}
 
 		// Configure logger
@@ -230,8 +231,31 @@ async function configureServerlessRunner(config: RunnerConfig): Promise<void> {
 				? config.autoConfigureServerless
 				: {};
 
+		// Make the request to fetch all datacenters
+		const dcsUrl = `${config.endpoint}/datacenters`;
+
+		logger().debug({
+			msg: "fetching datacenters",
+			url: dcsUrl,
+		});
+
+		const dcsResponse = await fetch(dcsUrl, {
+			headers: {
+				...(config.token ? { Authorization: `Bearer ${config.token}` } : {}),
+			},
+		});
+
+		if (!dcsResponse.ok) {
+			const errorText = await dcsResponse.text();
+			throw new Error(
+				`failed to configure serverless runner: ${dcsResponse.status} ${dcsResponse.statusText} - ${errorText}`,
+			);
+		}
+
+		let dcsRes = await dcsResponse.json() as { datacenters: { name: string }[] };
+
 		// Build the request body
-		const requestBody = {
+		const serverlessConfig = {
 			serverless: {
 				url:
 					customConfig.url ||
@@ -245,14 +269,15 @@ async function configureServerlessRunner(config: RunnerConfig): Promise<void> {
 					customConfig.slotsPerRunner ?? config.totalSlots ?? 1000,
 			},
 		};
+		const requestBody = Object.fromEntries(dcsRes.datacenters.map(dc => [dc.name, serverlessConfig]));
 
 		// Make the request to configure the serverless runner
 		const configUrl = `${config.endpoint}/runner-configs/${config.runnerName}?namespace=${config.namespace}`;
 
 		logger().debug({
 			msg: "configuring serverless runner",
 			url: configUrl,
-			config: requestBody.serverless,
+			config: serverlessConfig.serverless,
 		});
 
 		const response = await fetch(configUrl, {