chore(core): add token support

NathanFlurry · NathanFlurry · commit 8585e6a4962d · 2025-09-25T17:14:12.000-07:00
diff --git a/examples/next-js/package.json b/examples/next-js/package.json
@@ -21,4 +21,4 @@
     "@types/react": "^19",
     "@types/react-dom": "^19"
   }
-}
+}
diff --git a/packages/next-js/src/mod.ts b/packages/next-js/src/mod.ts
@@ -10,11 +10,14 @@ export const toNextHandler = (
 
 	const { fetch } = registry.startServerless(inputConfig);
 
-	const fetchWrapper = async (request: Request, { params }: { params: Promise<{ all: string[] }> }) => {
+	const fetchWrapper = async (
+		request: Request,
+		{ params }: { params: Promise<{ all: string[] }> },
+	) => {
 		const { all } = await params;
 
 		const newUrl = new URL(request.url);
-		newUrl.pathname = all.join('/');
+		newUrl.pathname = all.join("/");
 		const newReq = new Request(newUrl, request);
 
 		return await fetch(newReq);
diff --git a/packages/rivetkit/package.json b/packages/rivetkit/package.json
@@ -163,7 +163,7 @@
     "@bare-ts/lib": "~0.3.0",
     "@hono/standard-validator": "^0.1.3",
     "@hono/zod-openapi": "^0.19.10",
-    "@rivetkit/engine-runner": "https://pkg.pr.new/rivet-dev/engine/@rivetkit/engine-runner@a0639d0",
+    "@rivetkit/engine-runner": "https://pkg.pr.new/rivet-dev/engine/@rivetkit/engine-runner@0cfcf3d",
     "@rivetkit/fast-json-patch": "^3.1.2",
     "cbor-x": "^1.6.0",
     "hono": "^4.7.0",
diff --git a/packages/rivetkit/src/actor/driver.ts b/packages/rivetkit/src/actor/driver.ts
@@ -1,3 +1,4 @@
+import { Hono, MiddlewareHandler } from "hono";
 import type { CachedSerializer } from "@/actor/protocol/serde";
 import type { AnyClient } from "@/client/client";
 import type { ManagerDriver } from "@/manager/driver";
@@ -7,7 +8,6 @@ import type * as protocol from "@/schemas/client-protocol/mod";
 import type { AnyConn, ConnectionDriver } from "./connection";
 import type { GenericConnGlobalState } from "./generic-conn-driver";
 import type { AnyActorInstance } from "./instance";
-import { Hono, MiddlewareHandler } from "hono";
 
 export type ConnectionDriversMap = Record<ConnectionDriver, ConnDriver>;
 
diff --git a/packages/rivetkit/src/client/actor-conn.ts b/packages/rivetkit/src/client/actor-conn.ts
@@ -396,9 +396,9 @@ enc
 		logger().trace(
 			getEnvUniversal("_RIVETKIT_LOG_MESSAGE")
 				? {
-					msg: "parsed message",
-					message: jsonStringifyCompat(response).substring(0, 100) + "...",
-				}
+						msg: "parsed message",
+						message: jsonStringifyCompat(response).substring(0, 100) + "...",
+					}
 				: { msg: "parsed message" },
 		);
 
@@ -722,9 +722,9 @@ enc
 			logger().trace(
 				getEnvUniversal("_RIVETKIT_LOG_MESSAGE")
 					? {
-						msg: "sent http message",
-						message: `${jsonStringifyCompat(message).substring(0, 100)}...`,
-					}
+							msg: "sent http message",
+							message: `${jsonStringifyCompat(message).substring(0, 100)}...`,
+						}
 					: { msg: "sent http message" },
 			);
 
diff --git a/packages/rivetkit/src/client/config.ts b/packages/rivetkit/src/client/config.ts
@@ -16,8 +16,8 @@ export const ClientConfigSchema = z.object({
 
 	token: z
 		.string()
-		.nullable()
-		.default(() => getEnvUniversal("RIVET_TOKEN") ?? null),
+		.optional()
+		.transform((x) => x ?? getEnvUniversal("RIVET_TOKEN")),
 
 	headers: z.record(z.string()).optional().default({}),
 
diff --git a/packages/rivetkit/src/common/versioned-data.ts b/packages/rivetkit/src/common/versioned-data.ts
@@ -13,7 +13,7 @@ export interface VersionedDataConfig<T> {
 }
 
 export class VersionedDataHandler<T> {
-	constructor(private config: VersionedDataConfig<T>) { }
+	constructor(private config: VersionedDataConfig<T>) {}
 
 	serializeWithEmbeddedVersion(data: T): Uint8Array {
 		const versioned: VersionedData<Uint8Array> = {
diff --git a/packages/rivetkit/src/driver-test-suite/mod.ts b/packages/rivetkit/src/driver-test-suite/mod.ts
@@ -33,6 +33,7 @@ import { runRequestAccessTests } from "./tests/request-access";
 export interface SkipTests {
 	schedule?: boolean;
 	sleep?: boolean;
+	sse?: boolean;
 }
 
 export interface DriverTestConfig {
@@ -87,7 +88,10 @@ export function runDriverTests(
 			runActorDriverTests(driverTestConfig);
 			runManagerDriverTests(driverTestConfig);
 
-			for (const transport of ["websocket", "sse"] as Transport[]) {
+			const transports: Transport[] = driverTestConfig.skip?.sse
+				? ["websocket"]
+				: ["websocket", "sse"];
+			for (const transport of transports) {
 				describe(`transport (${transport})`, () => {
 					runActorConnTests({
 						...driverTestConfig,
diff --git a/packages/rivetkit/src/drivers/default.ts b/packages/rivetkit/src/drivers/default.ts
@@ -3,15 +3,12 @@ import { loggerWithoutContext } from "@/actor/log";
 import { createEngineDriver } from "@/drivers/engine/mod";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
 import type { DriverConfig, RunConfig } from "@/registry/run-config";
-import { getEnvUniversal } from "@/utils";
 
 /**
  * Chooses the appropriate driver based on the run configuration.
  */
 export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
-	const engineEndpoint = runConfig.endpoint ?? getEnvUniversal("RIVET_ENGINE");
-
-	if (engineEndpoint && runConfig.driver) {
+	if (runConfig.endpoint && runConfig.driver) {
 		throw new UserError(
 			"Cannot specify both 'engine' and 'driver' in configuration",
 		);
@@ -21,12 +18,16 @@ export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
 		return runConfig.driver;
 	}
 
-	if (engineEndpoint) {
+	if (runConfig.endpoint) {
 		loggerWithoutContext().debug({
 			msg: "using rivet engine driver",
-			endpoint: engineEndpoint,
+			endpoint: runConfig.endpoint,
+		});
+		// TODO: Add all properties from config
+		return createEngineDriver({
+			endpoint: runConfig.endpoint,
+			token: runConfig.token,
 		});
-		return createEngineDriver({ endpoint: engineEndpoint });
 	}
 
 	loggerWithoutContext().debug({ msg: "using default file system driver" });
diff --git a/packages/rivetkit/src/drivers/engine/actor-driver.ts b/packages/rivetkit/src/drivers/engine/actor-driver.ts
@@ -4,6 +4,8 @@ import type {
 } from "@rivetkit/engine-runner";
 import { Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
+import { Hono, MiddlewareHandler } from "hono";
+import { streamSSE } from "hono/streaming";
 import { WSContext } from "hono/ws";
 import invariant from "invariant";
 import { lookupInRegistry } from "@/actor/definition";
@@ -41,8 +43,6 @@ import { promiseWithResolvers } from "@/utils";
 import type { Config } from "./config";
 import { KEYS } from "./kv";
 import { logger } from "./log";
-import { Hono, MiddlewareHandler } from "hono";
-import { streamSSE } from "hono/streaming";
 
 interface ActorHandler {
 	actor?: AnyActorInstance;
@@ -86,12 +86,12 @@ 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,
 			runnerName: config.runnerName,
 			runnerKey: config.runnerKey,
-			token: config.token,
 			metadata: {
 				inspectorToken: this.#runConfig.inspector.token(),
 			},
@@ -304,15 +304,13 @@ export class EngineActorDriver implements ActorDriver {
 
 		const url = new URL(request.url);
 
+		// 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`);
 
-		// Parse configuration from Sec-WebSocket-Protocol header
-		const protocols = request.headers.get("sec-websocket-protocol");
 		let encodingRaw: string | undefined;
 		let connParamsRaw: string | undefined;
-		let token: string | undefined;
 
 		if (protocols) {
 			const protocolList = protocols.split(",").map((p) => p.trim());
@@ -323,8 +321,6 @@ export class EngineActorDriver implements ActorDriver {
 					connParamsRaw = decodeURIComponent(
 						protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
 					);
-				} else if (protocol.startsWith(WS_PROTOCOL_TOKEN)) {
-					token = protocol.substring(WS_PROTOCOL_TOKEN.length);
 				}
 			}
 		}
diff --git a/packages/rivetkit/src/drivers/engine/config.ts b/packages/rivetkit/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()
@@ -23,10 +27,6 @@ export const ConfigSchema = z
 			.default(
 				() => getEnvUniversal("RIVET_RUNNER_KEY") ?? crypto.randomUUID(),
 			),
-		token: z
-			.string()
-			.optional()
-			.transform((val) => val ?? getEnvUniversal("RIVET_TOKEN")),
 		totalSlots: z.number().default(100_000),
 	})
 	.default({});
diff --git a/packages/rivetkit/src/drivers/file-system/actor.ts b/packages/rivetkit/src/drivers/file-system/actor.ts
@@ -1,3 +1,5 @@
+import { Hono, MiddlewareHandler } from "hono";
+import { streamSSE } from "hono/streaming";
 import type { GenericConnGlobalState } from "@/actor/generic-conn-driver";
 import { loggerWithoutContext } from "@/actor/log";
 import type { AnyClient } from "@/client/client";
@@ -9,8 +11,6 @@ import type {
 import type { RegistryConfig, RunConfig } from "@/mod";
 import { bufferToArrayBuffer } from "@/utils";
 import type { FileSystemGlobalState } from "./global-state";
-import { Hono, MiddlewareHandler } from "hono";
-import { streamSSE } from "hono/streaming";
 
 export type ActorDriverContext = Record<never, never>;
 
diff --git a/packages/rivetkit/src/manager/router.ts b/packages/rivetkit/src/manager/router.ts
@@ -1,8 +1,14 @@
 import { createRoute, OpenAPIHono } from "@hono/zod-openapi";
 import * as cbor from "cbor-x";
-import { Hono, Context as HonoContext, Next, MiddlewareHandler } from "hono";
+import {
+	Hono,
+	Context as HonoContext,
+	type MiddlewareHandler,
+	Next,
+} from "hono";
 import { cors as corsMiddleware } from "hono/cors";
 import { createMiddleware } from "hono/factory";
+import { streamSSE } from "hono/streaming";
 import invariant from "invariant";
 import { z } from "zod";
 import { ActorNotFound, Unsupported } from "@/actor/errors";
@@ -21,7 +27,7 @@ import {
 	loggerMiddleware,
 } from "@/common/router";
 import { deconstructError, noopNext } from "@/common/utils";
-import { ActorDriver, HEADER_ACTOR_ID } from "@/driver-helpers/mod";
+import { type ActorDriver, HEADER_ACTOR_ID } from "@/driver-helpers/mod";
 import type {
 	TestInlineDriverCallRequest,
 	TestInlineDriverCallResponse,
@@ -48,7 +54,6 @@ import { stringifyError } from "@/utils";
 import type { ActorOutput, ManagerDriver } from "./driver";
 import { actorGateway, createTestWebSocketProxy } from "./gateway";
 import { logger } from "./log";
-import { streamSSE } from "hono/streaming";
 
 function buildOpenApiResponses<T>(schema: T) {
 	return {
diff --git a/packages/rivetkit/src/registry/mod.ts b/packages/rivetkit/src/registry/mod.ts
@@ -1,5 +1,6 @@
 import { type Client, createClientWithDriver } from "@/client/client";
 import { configureBaseLogger, configureDefaultLogger } from "@/common/log";
+import type { ActorDriver } from "@/driver-helpers/mod";
 import { chooseDefaultDriver } from "@/drivers/default";
 import {
 	configureInspectorAccessToken,
@@ -22,7 +23,6 @@ import {
 	RunConfigSchema,
 } from "./run-config";
 import { crossPlatformServe } from "./serve";
-import { ActorDriver } from "@/driver-helpers/mod";
 
 interface ServerOutput<A extends Registry<any>> {
 	/** Client to communicate with the actors. */
diff --git a/packages/rivetkit/src/remote-manager-driver/actor-http-client.ts b/packages/rivetkit/src/remote-manager-driver/actor-http-client.ts
@@ -25,10 +25,7 @@ export async function sendHttpRequestToActor(
 		actorId,
 	);
 
-	if (
-		actorRequest.method !== "GET" &&
-		actorRequest.method !== "HEAD"
-	) {
+	if (actorRequest.method !== "GET" && actorRequest.method !== "HEAD") {
 		if (actorRequest.bodyUsed) {
 			throw new Error("Request body has already been consumed");
 		}
@@ -37,16 +34,13 @@ export async function sendHttpRequestToActor(
 		// need to properly implement streaming bodies.
 		const reqBody = await actorRequest.arrayBuffer();
 
-		if (reqBody.byteLength != 0) {
+		if (reqBody.byteLength !== 0) {
 			bodyToSend = reqBody;
 
 			// If this is a streaming request, we need to convert the headers
 			// for the basic array buffer
 			guardHeaders.delete("transfer-encoding");
-			guardHeaders.set(
-				"content-length",
-				String(bodyToSend.byteLength),
-			);
+			guardHeaders.set("content-length", String(bodyToSend.byteLength));
 		}
 	}
 
diff --git a/packages/rivetkit/tests/driver-engine.test.ts b/packages/rivetkit/tests/driver-engine.test.ts
@@ -11,6 +11,7 @@ runDriverTests({
 	skip: {
 		// Skip tests that aren't applicable for engine-runner
 		schedule: true, // Scheduling handled by engine
+		sse: true,
 	},
 	async start() {
 		return await createTestRuntime(
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/vitest.base.ts b/vitest.base.ts

Original file line number	Diff line number	Diff line change
`@@ -21,4 +21,4 @@`
`21`	`21`	`"@types/react": "^19",`
`22`	`22`	`"@types/react-dom": "^19"`
`23`	`23`	`}`
`24`		`-}`
	`24`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ export interface VersionedDataConfig<T> {`
`13`	`13`	`}`
`14`	`14`
`15`	`15`	`export class VersionedDataHandler<T> {`
`16`		`- constructor(private config: VersionedDataConfig<T>) { }`
	`16`	`+ constructor(private config: VersionedDataConfig<T>) {}`
`17`	`17`
`18`	`18`	`serializeWithEmbeddedVersion(data: T): Uint8Array {`
`19`	`19`	`const versioned: VersionedData<Uint8Array> = {`