From 5650abed9d53b3840e4f4feff18ea4374952c2eb Mon Sep 17 00:00:00 2001
From: MasterPtato <maxpopkov.contact@gmail.com>
Date: Tue, 30 Sep 2025 14:08:53 -0700
Subject: [PATCH] feat: read total slots header for serverless

---
 packages/rivetkit/src/client/config.ts               |  2 ++
 packages/rivetkit/src/drivers/engine/actor-driver.ts |  4 ++--
 packages/rivetkit/src/manager/router.ts              | 10 ++++++----
 packages/rivetkit/src/registry/mod.ts                | 10 +++++++---
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/packages/rivetkit/src/client/config.ts b/packages/rivetkit/src/client/config.ts
index 8369152c8..73a5fa915 100644
--- a/packages/rivetkit/src/client/config.ts
+++ b/packages/rivetkit/src/client/config.ts
@@ -19,6 +19,8 @@ export const ClientConfigSchema = z.object({
 		.optional()
 		.transform((x) => x ?? getEnvUniversal("RIVET_TOKEN")),
 
+	totalSlots: z.number().optional(),
+
 	headers: z.record(z.string()).optional().default({}),
 
 	/** Endpoint to connect to the Rivet engine. Can be configured via RIVET_ENGINE env var. */
diff --git a/packages/rivetkit/src/drivers/engine/actor-driver.ts b/packages/rivetkit/src/drivers/engine/actor-driver.ts
index 755166901..c6db28bce 100644
--- a/packages/rivetkit/src/drivers/engine/actor-driver.ts
+++ b/packages/rivetkit/src/drivers/engine/actor-driver.ts
@@ -90,10 +90,10 @@ export class EngineActorDriver implements ActorDriver {
 		const runnerConfig: RunnerConfig = {
 			version: this.#version,
 			endpoint: config.endpoint,
-			token: config.token,
+			token: runConfig.token ?? config.token,
 			pegboardEndpoint: config.pegboardEndpoint,
 			namespace: config.namespace,
-			totalSlots: config.totalSlots,
+			totalSlots: runConfig.totalSlots ?? config.totalSlots,
 			runnerName: config.runnerName,
 			runnerKey: config.runnerKey,
 			metadata: {
diff --git a/packages/rivetkit/src/manager/router.ts b/packages/rivetkit/src/manager/router.ts
index 2ca344401..46125ce0a 100644
--- a/packages/rivetkit/src/manager/router.ts
+++ b/packages/rivetkit/src/manager/router.ts
@@ -105,7 +105,7 @@ export function createManagerRouter(
 }
 
 function addServerlessRoutes(
-	serverlessActorDriverBuilder: (token: string | undefined) => ActorDriver,
+	serverlessActorDriverBuilder: (token: string | undefined, totalSlots: number | undefined) => ActorDriver,
 	router: OpenAPIHono,
 	cors: MiddlewareHandler,
 ) {
@@ -118,9 +118,11 @@ function addServerlessRoutes(
 
 	// Serverless start endpoint
 	router.get("/start", cors, async (c) => {
-		const actorDriver = serverlessActorDriverBuilder(
-			c.req.header("x-rivet-token"),
-		);
+		let token = c.req.header("x-rivet-token");
+		let totalSlots: number | undefined = parseInt(c.req.header("x-rivetkit-total-slots") as any);
+		if (isNaN(totalSlots)) totalSlots = undefined;
+
+		const actorDriver = serverlessActorDriverBuilder(token, totalSlots);
 		invariant(
 			actorDriver.serverlessHandleStart,
 			"missing serverlessHandleStart on ActorDriver",
diff --git a/packages/rivetkit/src/registry/mod.ts b/packages/rivetkit/src/registry/mod.ts
index 755395025..64b4e8d08 100644
--- a/packages/rivetkit/src/registry/mod.ts
+++ b/packages/rivetkit/src/registry/mod.ts
@@ -216,10 +216,14 @@ export class Registry<A extends RegistryActors> {
 		}
 
 		let serverlessActorDriverBuilder:
-			| ((token?: string) => ActorDriver)
-			| undefined = (token: string | undefined) => {
-			// Override config token if one was set
+			| ((token?: string, totalSlots?: number) => ActorDriver)
+			| undefined = (
+			token: string | undefined,
+			totalSlots: number | undefined,
+		) => {
+			// Override config
 			if (token) config.token = token;
+			if (totalSlots) config.totalSlots = totalSlots;
 
 			return driver.actor(this.#config, config, managerDriver, client);
 		};