feat: read total slots header for serverless (#1319)

Author: MasterPtato

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. */

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: {

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",

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);
 		};