chore(core): clean up token configs (#1298)

Author: NathanFlurry

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

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({}),
 

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,

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

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

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

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

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(

pnpm-lock.yaml

@@ -2,8 +2,8 @@ import type { ViteUserConfig } from "vitest/config";
 
 export default {
 	test: {
-		testTimeout: 5_000,
-		hookTimeout: 5_000,
+		testTimeout: 10_000,
+		hookTimeout: 10_000,
 		// Enable parallelism
 		sequence: {
 			// TODO: This breaks fake timers, unsure how to make tests run in parallel within the same file