chore(rivetkit): auto-generate inspector token for engine driver

Author: NathanFlurry

rivetkit-typescript/packages/cloudflare-workers/src/actor-driver.ts

@@ -6,10 +6,11 @@ import type {
 } from "rivetkit";
 import { lookupInRegistry } from "rivetkit";
 import type { Client } from "rivetkit/client";
-import type {
-	ActorDriver,
-	AnyActorInstance,
-	ManagerDriver,
+import {
+	type ActorDriver,
+	type AnyActorInstance,
+	generateRandomString,
+	type ManagerDriver,
 } from "rivetkit/driver-helpers";
 import { promiseWithResolvers } from "rivetkit/utils";
 import { KEYS } from "./actor-handler-do";
@@ -154,6 +155,11 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 	async getDatabase(actorId: string): Promise<unknown | undefined> {
 		return this.#getDOCtx(actorId).storage.sql;
 	}
+
+	getInspectorToken(_actorId: string): string | undefined {
+		// TODO: Inspector currently disabled. Need to implement a way to store the token in CF KV simplar to the Engine driver.
+		return generateRandomString();
+	}
 }
 
 export function createCloudflareActorsActorDriverBuilder(

rivetkit-typescript/packages/rivetkit/src/actor/driver.ts

@@ -53,4 +53,11 @@ export interface ActorDriver {
 
 	/** Extra properties to add to logs for each actor. */
 	getExtraActorLogParams?(): Record<string, string>;
+
+	/**
+	 * Get the inspector token for an actor.
+	 *
+	 * If returns undefined, auth is disabled.
+	 **/
+	getInspectorToken(actorId: string): string | undefined;
 }

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

@@ -19,6 +19,7 @@ import {
 	handleWebSocketConnect,
 	truncateRawWebSocketPathPrefix,
 } from "@/actor/router-endpoints";
+import { generateSecureToken } from "@/actor/utils";
 import type { Client } from "@/client/client";
 import {
 	PATH_CONNECT_WEBSOCKET,
@@ -60,6 +61,7 @@ interface ActorHandler {
 	actor?: AnyActorInstance;
 	actorStartPromise?: ReturnType<typeof promiseWithResolvers<void>>;
 	persistedData?: Uint8Array;
+	inspectorToken?: string;
 }
 
 export type DriverContext = {};
@@ -438,9 +440,11 @@ export class EngineActorDriver implements ActorDriver {
 			this.#actors.set(actorId, handler);
 
 			// Load persisted data from storage
-			const [persistedValue] = await this.#runner.kvGet(actorId, [
-				KEYS.PERSIST_DATA,
-			]);
+			const [persistedValue, inspectorTokenValue] =
+				await this.#runner.kvGet(actorId, [
+					KEYS.PERSIST_DATA,
+					KEYS.INSPECTOR.TOKEN,
+				]);
 
 			handler.persistedData =
 				persistedValue !== null
@@ -453,6 +457,30 @@ export class EngineActorDriver implements ActorDriver {
 				dataSize: handler.persistedData?.byteLength,
 				wasInStorage: persistedValue !== null,
 			});
+
+			// Get or create inspector token
+			if (inspectorTokenValue !== null) {
+				// Token exists, decode it
+				const decoder = new TextDecoder();
+				handler.inspectorToken = decoder.decode(inspectorTokenValue);
+				logger().debug({
+					msg: "loaded existing inspector token",
+					actorId,
+				});
+			} else {
+				// Create new token
+				handler.inspectorToken = generateSecureToken();
+				const tokenBytes = new TextEncoder().encode(
+					handler.inspectorToken,
+				);
+				await this.#runner.kvPut(actorId, [
+					[KEYS.INSPECTOR.TOKEN, tokenBytes],
+				]);
+				logger().debug({
+					msg: "created new inspector token",
+					actorId,
+				});
+			}
 		}
 
 		const name = actorConfig.name as string;
@@ -909,4 +937,14 @@ export class EngineActorDriver implements ActorDriver {
 	getExtraActorLogParams(): Record<string, string> {
 		return { runnerId: this.#runner.runnerId ?? "-" };
 	}
+
+	getInspectorToken(actorId: string): string | undefined {
+		const handler = this.#actors.get(actorId);
+		invariant(handler, "handler does not exist in getInspectorToken");
+		invariant(
+			handler.inspectorToken,
+			"handler.inspectorToken does not exist in getInspectorToken",
+		);
+		return handler.inspectorToken;
+	}
 }

rivetkit-typescript/packages/rivetkit/src/drivers/engine/kv.ts

@@ -1,3 +1,6 @@
 export const KEYS = {
 	PERSIST_DATA: Uint8Array.from([1]),
+	INSPECTOR: {
+		TOKEN: Uint8Array.from([2]),
+	},
 };

rivetkit-typescript/packages/rivetkit/src/drivers/file-system/actor.ts

@@ -83,4 +83,8 @@ export class FileSystemActorDriver implements ActorDriver {
 		// Spawns the sleepActor promise
 		this.#state.sleepActor(actorId);
 	}
+
+	getInspectorToken(_actorId: string): string | undefined {
+		return undefined;
+	}
 }

rivetkit-typescript/packages/rivetkit/src/inspector/utils.ts

@@ -77,6 +77,7 @@ export function getInspectorUrl(runConfig: RunnerConfigInput | undefined) {
 
 export const isInspectorEnabled = (
 	runConfig: RunConfig,
+	// TODO(kacper): Remove context in favor of using the gateway, so only context is the actor
 	context: "actor" | "manager",
 ) => {
 	if (typeof runConfig.inspector?.enabled === "boolean") {

rivetkit-typescript/packages/rivetkit/src/manager/driver.ts

@@ -48,11 +48,13 @@ export interface ManagerDriver {
 		router: Hono,
 	) => void;
 
+	// TODO(kacper): Remove this in favor of standard manager API
 	/**
 	 * @internal
 	 */
 	readonly inspector?: ManagerInspector;
 
+	// TODO(kacper): Remove this in favor of ActorDriver.getinspectorToken
 	/**
 	 * Get or create the inspector access token.
 	 * @experimental

rivetkit-typescript/packages/rivetkit/src/manager/router.ts

@@ -222,6 +222,7 @@ function addManagerRoutes(
 	managerDriver: ManagerDriver,
 	router: OpenAPIHono,
 ) {
+	// TODO(kacper): Remove this in favor of standard manager API
 	// Inspector
 	if (isInspectorEnabled(runConfig, "manager")) {
 		if (!managerDriver.inspector) {