rivetkit 2.0.6 → 2.0.7

package/src/drivers/engine/actor-driver.ts

@@ -4,13 +4,11 @@ import type {
 } from "@rivetkit/engine-runner";
 import { Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
+import type { Context as HonoContext } from "hono";
+import { streamSSE } from "hono/streaming";
 import { WSContext } from "hono/ws";
 import invariant from "invariant";
 import { lookupInRegistry } from "@/actor/definition";
-import {
-	createGenericConnDrivers,
-	GenericConnGlobalState,
-} from "@/actor/generic-conn-driver";
 import { deserializeActorKey } from "@/actor/keys";
 import { EncodingSchema } from "@/actor/protocol/serde";
 import { type ActorRouter, createActorRouter } from "@/actor/router";
@@ -22,6 +20,9 @@ import type { Client } from "@/client/client";
 import {
 	PATH_CONNECT_WEBSOCKET,
 	PATH_RAW_WEBSOCKET_PREFIX,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_TOKEN,
 } from "@/common/actor-router-consts";
 import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { getLogger } from "@/common/log";
@@ -29,22 +30,19 @@ import type { UniversalWebSocket } from "@/common/websocket-interface";
 import {
 	type ActorDriver,
 	type AnyActorInstance,
-	HEADER_AUTH_DATA,
-	HEADER_CONN_PARAMS,
-	HEADER_ENCODING,
 	type ManagerDriver,
 	serializeEmptyPersistData,
 } from "@/driver-helpers/mod";
 import type { RegistryConfig } from "@/registry/config";
 import type { RunConfig } from "@/registry/run-config";
+import { promiseWithResolvers } from "@/utils";
 import type { Config } from "./config";
 import { KEYS } from "./kv";
 import { logger } from "./log";
 
 interface ActorHandler {
 	actor?: AnyActorInstance;
-	actorStartPromise?: PromiseWithResolvers<void>;
-	genericConnGlobalState: GenericConnGlobalState;
+	actorStartPromise?: ReturnType<typeof promiseWithResolvers<void>>;
 	persistedData?: Uint8Array;
 }
 
@@ -61,6 +59,9 @@ export class EngineActorDriver implements ActorDriver {
 	#actorRouter: ActorRouter;
 	#version: number = 1; // Version for the runner protocol
 
+	#runnerStarted: PromiseWithResolvers<undefined> = Promise.withResolvers();
+	#runnerStopped: PromiseWithResolvers<undefined> = Promise.withResolvers();
+
 	constructor(
 		registryConfig: RegistryConfig,
 		runConfig: RunConfig,
@@ -80,6 +81,7 @@ export class EngineActorDriver implements ActorDriver {
 		const runnerConfig: RunnerConfig = {
 			version: this.#version,
 			endpoint: config.endpoint,
+			token: config.token,
 			pegboardEndpoint: config.pegboardEndpoint,
 			namespace: config.namespace,
 			totalSlots: config.totalSlots,
@@ -108,6 +110,8 @@ export class EngineActorDriver implements ActorDriver {
 						runnerName: this.#config.runnerName,
 					});
 				}
+
+				this.#runnerStarted.resolve(undefined);
 			},
 			onDisconnected: () => {
 				logger().warn({
@@ -117,7 +121,9 @@ export class EngineActorDriver implements ActorDriver {
 				});
 				hasDisconnected = true;
 			},
-			onShutdown: () => {},
+			onShutdown: () => {
+				this.#runnerStopped.resolve(undefined);
+			},
 			fetch: this.#runnerFetch.bind(this),
 			websocket: this.#runnerWebSocket.bind(this),
 			onActorStart: this.#runnerOnActorStart.bind(this),
@@ -151,14 +157,6 @@ export class EngineActorDriver implements ActorDriver {
 		return handler.actor;
 	}
 
-	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
-		const handler = this.#actors.get(actorId);
-		if (!handler) {
-			throw new Error(`Actor ${actorId} not loaded`);
-		}
-		return handler.genericConnGlobalState;
-	}
-
 	getContext(actorId: string): DriverContext {
 		return {};
 	}
@@ -221,8 +219,7 @@ export class EngineActorDriver implements ActorDriver {
 		let handler = this.#actors.get(actorId);
 		if (!handler) {
 			handler = {
-				genericConnGlobalState: new GenericConnGlobalState(),
-				actorStartPromise: Promise.withResolvers(),
+				actorStartPromise: promiseWithResolvers(),
 				persistedData: serializeEmptyPersistData(input),
 			};
 			this.#actors.set(actorId, handler);
@@ -240,11 +237,7 @@ export class EngineActorDriver implements ActorDriver {
 		handler.actor = definition.instantiate();
 
 		// Start actor
-		const connDrivers = createGenericConnDrivers(
-			handler.genericConnGlobalState,
-		);
 		await handler.actor.start(
-			connDrivers,
 			this,
 			this.#inlineClient,
 			actorId,
@@ -293,14 +286,29 @@ export class EngineActorDriver implements ActorDriver {
 
 		const url = new URL(request.url);
 
-		// Parse headers
-		const encodingRaw = request.headers.get(HEADER_ENCODING);
-		const connParamsRaw = request.headers.get(HEADER_CONN_PARAMS);
-		const authDataRaw = request.headers.get(HEADER_AUTH_DATA);
+		// Parse configuration from Sec-WebSocket-Protocol header
+		const protocols = request.headers.get("sec-websocket-protocol");
+		if (protocols === null)
+			throw new Error(`Missing sec-websocket-protocol header`);
+
+		let encodingRaw: string | undefined;
+		let connParamsRaw: string | undefined;
+
+		if (protocols) {
+			const protocolList = protocols.split(",").map((p) => p.trim());
+			for (const protocol of protocolList) {
+				if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
+					encodingRaw = protocol.substring(WS_PROTOCOL_ENCODING.length);
+				} else if (protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)) {
+					connParamsRaw = decodeURIComponent(
+						protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
+					);
+				}
+			}
+		}
 
 		const encoding = EncodingSchema.parse(encodingRaw);
 		const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;
-		const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
 
 		// Fetch WS handler
 		//
@@ -314,7 +322,9 @@ export class EngineActorDriver implements ActorDriver {
 				actorId,
 				encoding,
 				connParams,
-				authData,
+				// Extract connId and connToken from protocols if needed
+				undefined,
+				undefined,
 			);
 		} else if (url.pathname.startsWith(PATH_RAW_WEBSOCKET_PREFIX)) {
 			wsHandlerPromise = handleRawWebSocketHandler(
@@ -322,7 +332,6 @@ export class EngineActorDriver implements ActorDriver {
 				url.pathname + url.search,
 				this,
 				actorId,
-				authData,
 			);
 		} else {
 			throw new Error(`Unreachable path: ${url.pathname}`);
@@ -367,4 +376,17 @@ export class EngineActorDriver implements ActorDriver {
 		logger().info({ msg: "stopping engine actor driver" });
 		await this.#runner.shutdown(immediate);
 	}
+
+	async serverlessHandleStart(c: HonoContext): Promise<Response> {
+		await this.#runnerStarted.promise;
+
+		return streamSSE(c, async (stream) => {
+			// Runner id should be set if the runner started
+			const runnerId = this.#runner.runnerId;
+			invariant(runnerId, "runnerId not set");
+			stream.writeSSE({ data: runnerId });
+
+			return this.#runnerStopped.promise;
+		});
+	}
 }

package/src/drivers/engine/config.ts

@@ -10,6 +10,10 @@ export const ConfigSchema = z
 			.default(
 				() => getEnvUniversal("RIVET_ENGINE") ?? "http://localhost:6420",
 			),
+		token: z
+			.string()
+			.optional()
+			.transform((val) => val ?? getEnvUniversal("RIVET_TOKEN")),
 		pegboardEndpoint: z.string().optional(),
 		namespace: z
 			.string()

package/src/drivers/file-system/actor.ts

@@ -1,5 +1,3 @@
-import type { GenericConnGlobalState } from "@/actor/generic-conn-driver";
-import { loggerWithoutContext } from "@/actor/log";
 import type { AnyClient } from "@/client/client";
 import type {
 	ActorDriver,
@@ -46,10 +44,6 @@ export class FileSystemActorDriver implements ActorDriver {
 		);
 	}
 
-	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
-		return this.#state.getActorOrError(actorId).genericConnGlobalState;
-	}
-
 	/**
 	 * Get the current storage directory path
 	 */

package/src/drivers/file-system/global-state.ts

@@ -5,10 +5,6 @@ import * as path from "node:path";
 import invariant from "invariant";
 import { lookupInRegistry } from "@/actor/definition";
 import { ActorAlreadyExists } from "@/actor/errors";
-import {
-	createGenericConnDrivers,
-	GenericConnGlobalState,
-} from "@/actor/generic-conn-driver";
 import type { AnyActorInstance } from "@/actor/instance";
 import type { ActorKey } from "@/actor/mod";
 import { generateRandomString } from "@/actor/utils";
@@ -27,7 +23,7 @@ import {
 import {
 	bufferToArrayBuffer,
 	type LongTimeoutHandle,
-	SinglePromiseQueue,
+	promiseWithResolvers,
 	setLongTimeout,
 	stringifyError,
 } from "@/utils";
@@ -49,9 +45,7 @@ interface ActorEntry {
 
 	actor?: AnyActorInstance;
 	/** Promise for starting the actor. */
-	startPromise?: PromiseWithResolvers<void>;
-
-	genericConnGlobalState: GenericConnGlobalState;
+	startPromise?: ReturnType<typeof promiseWithResolvers<void>>;
 
 	alarmTimeout?: LongTimeoutHandle;
 	/** The timestamp currently scheduled for this actor's alarm (ms since epoch). */
@@ -188,7 +182,6 @@ export class FileSystemGlobalState {
 
 		entry = {
 			id: actorId,
-			genericConnGlobalState: new GenericConnGlobalState(),
 			removed: false,
 		};
 		this.#actors.set(actorId, entry);
@@ -469,7 +462,7 @@ export class FileSystemGlobalState {
 		}
 
 		// Create start promise
-		entry.startPromise = Promise.withResolvers();
+		entry.startPromise = promiseWithResolvers();
 
 		try {
 			// Create actor
@@ -477,11 +470,7 @@ export class FileSystemGlobalState {
 			entry.actor = definition.instantiate();
 
 			// Start actor
-			const connDrivers = createGenericConnDrivers(
-				entry.genericConnGlobalState,
-			);
 			await entry.actor.start(
-				connDrivers,
 				actorDriver,
 				inlineClient,
 				actorId,

package/src/drivers/file-system/manager.ts

@@ -57,10 +57,6 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		this.#driverConfig = driverConfig;
 
 		if (runConfig.inspector.enabled) {
-			if (!this.#runConfig.inspector.token()) {
-				this.#runConfig.inspector.token = () =>
-					this.#state.getOrCreateInspectorAccessToken();
-			}
 			const startedAt = new Date().toISOString();
 			function transformActor(actorState: schema.ActorState): Actor {
 				return {
@@ -141,6 +137,8 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<UniversalWebSocket> {
 		// Handle raw WebSocket paths
 		const pathOnly = path.split("?")[0];
@@ -154,7 +152,8 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				params,
-				undefined,
+				connId,
+				connToken,
 			);
 			return new InlineWebSocketAdapter2(wsHandler);
 		} else if (
@@ -168,7 +167,6 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				path,
 				this.#actorDriver,
 				actorId,
-				undefined,
 			);
 			return new InlineWebSocketAdapter2(wsHandler);
 		} else {
@@ -192,6 +190,8 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		connParams: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<Response> {
 		const upgradeWebSocket = this.#runConfig.getUpgradeWebSocket?.();
 		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
@@ -208,7 +208,8 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				connParams,
-				undefined,
+				connId,
+				connToken,
 			);
 			return upgradeWebSocket(() => wsHandler)(c, noopNext());
 		} else if (
@@ -222,7 +223,6 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				path,
 				this.#actorDriver,
 				actorId,
-				undefined,
 			);
 			return upgradeWebSocket(() => wsHandler)(c, noopNext());
 		} else {
@@ -321,4 +321,8 @@ export class FileSystemManagerDriver implements ManagerDriver {
 			data: this.#state.storagePath,
 		};
 	}
+
+	getOrCreateInspectorAccessToken() {
+		return this.#state.getOrCreateInspectorAccessToken();
+	}
 }

package/src/inspector/actor.ts

@@ -8,6 +8,7 @@ import type {
 	AnyDatabaseProvider,
 	InferDatabaseClient,
 } from "@/actor/database";
+import { promiseWithResolvers } from "@/utils";
 import {
 	ColumnsSchema,
 	type Connection,
@@ -99,7 +100,7 @@ export function createActorInspectorRouter() {
 						});
 					});
 
-					const { promise } = Promise.withResolvers<void>();
+					const { promise } = promiseWithResolvers<void>();
 
 					return promise;
 				},
@@ -128,7 +129,7 @@ export function createActorInspectorRouter() {
 						});
 					});
 
-					const { promise } = Promise.withResolvers<void>();
+					const { promise } = promiseWithResolvers<void>();
 
 					return promise;
 				},
@@ -159,7 +160,7 @@ export function createActorInspectorRouter() {
 						});
 					});
 
-					const { promise } = Promise.withResolvers<void>();
+					const { promise } = promiseWithResolvers<void>();
 
 					return promise;
 				},

package/src/inspector/config.ts

@@ -56,7 +56,16 @@ const defaultCors: CorsOptions = {
 
 export const InspectorConfigSchema = z
 	.object({
-		enabled: z.boolean().optional().default(defaultEnabled),
+		enabled: z
+			.boolean()
+			.or(
+				z.object({
+					actor: z.boolean().optional().default(true),
+					manager: z.boolean().optional().default(true),
+				}),
+			)
+			.optional()
+			.default(defaultEnabled),
 		/** CORS configuration for the router. Uses Hono's CORS middleware options. */
 		cors: z
 			.custom<CorsOptions>()

package/src/inspector/mod.ts

@@ -1,2 +1,3 @@
 export * from "./protocol/common";
 export * from "./protocol/mod";
+export * from "./utils";

package/src/inspector/utils.ts

@@ -1,5 +1,6 @@
 import crypto from "node:crypto";
 import { createMiddleware } from "hono/factory";
+import type { ManagerDriver } from "@/driver-helpers/mod";
 import type { RunConfig } from "@/mod";
 import type { RunConfigInput } from "@/registry/run-config";
 import { inspectorLogger } from "./log";
@@ -28,10 +29,6 @@ export function compareSecrets(providedSecret: string, validSecret: string) {
 
 export const secureInspector = (runConfig: RunConfig) =>
 	createMiddleware(async (c, next) => {
-		if (!runConfig.inspector.enabled) {
-			return c.text("Inspector is not enabled", 503);
-		}
-
 		const userToken = c.req.header("Authorization")?.replace("Bearer ", "");
 		if (!userToken) {
 			return c.text("Unauthorized", 401);
@@ -74,3 +71,25 @@ export function getInspectorUrl(runConfig: RunConfigInput | undefined) {
 
 	return url.href;
 }
+
+export const isInspectorEnabled = (
+	runConfig: RunConfig,
+	context: "actor" | "manager",
+) => {
+	if (typeof runConfig.inspector?.enabled === "boolean") {
+		return runConfig.inspector.enabled;
+	} else if (typeof runConfig.inspector?.enabled === "object") {
+		return runConfig.inspector.enabled[context];
+	}
+	return false;
+};
+
+export const configureInspectorAccessToken = (
+	runConfig: RunConfig,
+	managerDriver: ManagerDriver,
+) => {
+	if (!runConfig.inspector?.token()) {
+		const token = managerDriver.getOrCreateInspectorAccessToken();
+		runConfig.inspector.token = () => token;
+	}
+};

package/src/manager/driver.ts

@@ -21,6 +21,8 @@ export interface ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<UniversalWebSocket>;
 	proxyRequest(
 		c: HonoContext,
@@ -33,7 +35,8 @@ export interface ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
-		authData: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<Response>;
 
 	displayInformation(): ManagerDisplayInformation;
@@ -46,6 +49,13 @@ export interface ManagerDriver {
 	 * @internal
 	 */
 	readonly inspector?: ManagerInspector;
+
+	/**
+	 * Get or create the inspector access token.
+	 * @internal
+	 * @returns creates or returns existing inspector access token
+	 */
+	getOrCreateInspectorAccessToken: () => string;
 }
 
 export interface ManagerDisplayInformation {