feat: add onAuth lifecycle hooks

Author: NathanFlurry

CLAUDE.md

@@ -108,3 +108,7 @@ This ensures imports resolve correctly across different build environments and p
 - Run `yarn check-types` regularly during development to catch type errors early. Prefer `yarn check-types` instead of `yarn build`.
 - Use `tsx` CLI to execute TypeScript scripts directly (e.g., `tsx script.ts` instead of `node script.js`).
 - Do not auto-commit changes
+
+## Test Guidelines
+
+- Do not check if errors are an instanceOf WorkerError in tests. Many error types do not have the same prototype chain when sent over the network, but still have the same properties so you can safely cast with `as`.

docs/workers/quickstart.mdx

@@ -7,10 +7,11 @@ description: Start building awesome documentation in under 5 minutes
 ```ts registry.ts
 import { setup } from "rivetkit";
 import { worker } from "rivetkit/worker";
-import { workflow } from "rivetkit/workflow";
-import { realtime } from "rivetkit/realtime";
 
 const counter = worker({
+	onAuth: async () => {
+		// Allow public access
+	},
 	state: {
 		count: 0,
 	},
@@ -36,18 +37,24 @@ export type Registry = typeof registry;
 ```
 
 ```ts server.ts
+// With router
 import { registry } from "./registry";
 
-const registry = new Hono();
-app.route("/registry", registry.handler);
+const app = new Hono();
+app.route("/registry", c => registry.handler(c.req.raw));
+serve(app);
+
+// Without router
+import { serve } from "@rivetkit/node";
+
 serve(registry);
 ```
 
 ```ts client.ts
 import { createClient } from "rivetkit/client";
 import type { Registry } from "./registry";
 
-const client = createClient<Registry>("http://localhost:8080/registry");
+const client = createClient<Registry>("http://localhost:8080");
 ```
 
 <Steps>

packages/core/fixtures/driver-test-suite/auth.ts

@@ -0,0 +1,105 @@
+import { worker, UserError } from "rivetkit";
+
+// Basic auth worker - requires API key
+export const authWorker = worker({
+	state: { requests: 0 },
+	onAuth: (opts) => {
+		const { req, intents, params } = opts;
+		const apiKey = (params as any)?.apiKey;
+		if (!apiKey) {
+			throw new UserError("API key required", { code: "missing_auth" });
+		}
+		
+		if (apiKey !== "valid-api-key") {
+			throw new UserError("Invalid API key", { code: "invalid_auth" });
+		}
+		
+		return { userId: "user123", token: apiKey };
+	},
+	actions: {
+		getRequests: (c) => {
+			c.state.requests++;
+			return c.state.requests;
+		},
+		getUserAuth: (c) => c.conn.auth,
+	},
+});
+
+// Intent-specific auth worker - checks different permissions for different intents
+export const intentAuthWorker = worker({
+	state: { value: 0 },
+	onAuth: (opts) => {
+		const { req, intents, params } = opts;
+		console.log('intents', intents, params);
+		const role = (params as any)?.role;
+		
+		if (intents.has("create") && role !== "admin") {
+			throw new UserError("Admin role required for create operations", { code: "insufficient_permissions" });
+		}
+		
+		if (intents.has("action") && !["admin", "user"].includes(role || "")) {
+			throw new UserError("User or admin role required for actions", { code: "insufficient_permissions" });
+		}
+		
+		return { role, timestamp: Date.now() };
+	},
+	actions: {
+		getValue: (c) => c.state.value,
+		setValue: (c, value: number) => {
+			c.state.value = value;
+			return value;
+		},
+		getAuth: (c) => c.conn.auth,
+	},
+});
+
+// Public worker - empty onAuth to allow public access
+export const publicWorker = worker({
+	state: { visitors: 0 },
+	onAuth: () => {
+		return null; // Allow public access
+	},
+	actions: {
+		visit: (c) => {
+			c.state.visitors++;
+			return c.state.visitors;
+		},
+	},
+});
+
+// No auth worker - should fail when accessed publicly (no onAuth defined)
+export const noAuthWorker = worker({
+	state: { value: 42 },
+	actions: {
+		getValue: (c) => c.state.value,
+	},
+});
+
+// Async auth worker - tests promise-based authentication
+export const asyncAuthWorker = worker({
+	state: { count: 0 },
+	onAuth: async (opts) => {
+		const { req, intents, params } = opts;
+		// Simulate async auth check (e.g., database lookup)
+		await new Promise(resolve => setTimeout(resolve, 10));
+		
+		const token = (params as any)?.token;
+		if (!token) {
+			throw new UserError("Token required", { code: "missing_token" });
+		}
+		
+		// Simulate token validation
+		if (token === "invalid") {
+			throw new UserError("Token is invalid", { code: "invalid_token" });
+		}
+		
+		return { userId: `user-${token}`, validated: true };
+	},
+	actions: {
+		increment: (c) => {
+			c.state.count++;
+			return c.state.count;
+		},
+		getAuthData: (c) => c.conn.auth,
+	},
+});

packages/core/fixtures/driver-test-suite/registry.ts

@@ -27,6 +27,13 @@ import {
 	uniqueVarWorker,
 	driverCtxWorker,
 } from "./vars";
+import {
+	authWorker,
+	intentAuthWorker,
+	publicWorker,
+	noAuthWorker,
+	asyncAuthWorker,
+} from "./auth";
 
 // Consolidated setup with all workers
 export const registry = setup({
@@ -63,6 +70,12 @@ export const registry = setup({
 		dynamicVarWorker,
 		uniqueVarWorker,
 		driverCtxWorker,
+		// From auth.ts
+		authWorker,
+		intentAuthWorker,
+		publicWorker,
+		noAuthWorker,
+		asyncAuthWorker,
 	},
 });
 

packages/core/src/client/client.ts

@@ -172,6 +172,7 @@ export interface ClientDriver {
 		c: HonoContext | undefined,
 		workerQuery: WorkerQuery,
 		encodingKind: Encoding,
+		params: unknown,
 	): Promise<string>;
 	connectWebSocket(
 		c: HonoContext | undefined,
@@ -364,6 +365,7 @@ export class ClientRaw {
 			undefined,
 			createQuery,
 			this.#encodingKind,
+			opts?.params,
 		);
 		logger().debug("created worker with ID", {
 			name,
@@ -481,11 +483,9 @@ export function createClientWithDriver<A extends Registry<any>>(
 						key?: string | string[],
 						opts?: GetOptions,
 					): WorkerHandle<ExtractWorkersFromRegistry<A>[typeof prop]> => {
-						return target.getOrCreate<ExtractWorkersFromRegistry<A>[typeof prop]>(
-							prop,
-							key,
-							opts,
-						);
+						return target.getOrCreate<
+							ExtractWorkersFromRegistry<A>[typeof prop]
+						>(prop, key, opts);
 					},
 					getForId: (
 						workerId: string,
@@ -500,12 +500,12 @@ export function createClientWithDriver<A extends Registry<any>>(
 					create: async (
 						key: string | string[],
 						opts: CreateOptions = {},
-					): Promise<WorkerHandle<ExtractWorkersFromRegistry<A>[typeof prop]>> => {
-						return await target.create<ExtractWorkersFromRegistry<A>[typeof prop]>(
-							prop,
-							key,
-							opts,
-						);
+					): Promise<
+						WorkerHandle<ExtractWorkersFromRegistry<A>[typeof prop]>
+					> => {
+						return await target.create<
+							ExtractWorkersFromRegistry<A>[typeof prop]
+						>(prop, key, opts);
 					},
 				} as WorkerAccessor<ExtractWorkersFromRegistry<A>[typeof prop]>;
 			}

packages/core/src/client/http-client-driver.ts

@@ -82,6 +82,7 @@ export function createHttpClientDriver(managerEndpoint: string): ClientDriver {
 			_c: HonoContext | undefined,
 			workerQuery: WorkerQuery,
 			encodingKind: Encoding,
+			params: unknown,
 		): Promise<string> => {
 			logger().debug("resolving worker ID", { query: workerQuery });
 
@@ -95,6 +96,9 @@ export function createHttpClientDriver(managerEndpoint: string): ClientDriver {
 					headers: {
 						[HEADER_ENCODING]: encodingKind,
 						[HEADER_WORKER_QUERY]: JSON.stringify(workerQuery),
+						...(params !== undefined
+							? { [HEADER_CONN_PARAMS]: JSON.stringify(params) }
+							: {}),
 					},
 					body: {},
 					encoding: encodingKind,
@@ -122,14 +126,20 @@ export function createHttpClientDriver(managerEndpoint: string): ClientDriver {
 		): Promise<WebSocket> => {
 			const { WebSocket } = await dynamicImports;
 
-			const workerQueryStr = encodeURIComponent(JSON.stringify(workerQuery));
 			const endpoint = managerEndpoint
 				.replace(/^http:/, "ws:")
 				.replace(/^https:/, "wss:");
-			const url = `${endpoint}/workers/connect/websocket?encoding=${encodingKind}&query=${workerQueryStr}`;
+			const url = `${endpoint}/workers/connect/websocket`;
+
+			// Pass sensitive data via protocol
+			const protocol = [
+				`query.${btoa(JSON.stringify(workerQuery))}`,
+				"encoding.encodingKind",
+			];
+			if (params) protocol.push(`conn_params.${btoa(JSON.stringify(params))}`);
 
 			logger().debug("connecting to websocket", { url });
-			const ws = new WebSocket(url);
+			const ws = new WebSocket(url, protocol);
 			if (encodingKind === "cbor") {
 				ws.binaryType = "arraybuffer";
 			} else if (encodingKind === "json") {

packages/core/src/client/worker-handle.ts

@@ -109,6 +109,7 @@ export class WorkerHandleRaw {
 				undefined,
 				this.#workerQuery,
 				this.#encodingKind,
+				this.#params,
 			);
 			this.#workerQuery = { getForId: { workerId } };
 			return workerId;

packages/core/src/driver-test-suite/mod.ts

@@ -21,6 +21,7 @@ import { runWorkerVarsTests } from "./tests/worker-vars";
 import { runWorkerConnStateTests } from "./tests/worker-conn-state";
 import { runWorkerMetadataTests } from "./tests/worker-metadata";
 import { runWorkerErrorHandlingTests } from "./tests/worker-error-handling";
+import { runWorkerAuthTests } from "./tests/worker-auth";
 
 export interface DriverTestConfig {
 	/** Deploys an registry and returns the connection endpoint. */
@@ -90,6 +91,8 @@ export function runDriverTests(
 			runWorkerMetadataTests(driverTestConfig);
 
 			runWorkerErrorHandlingTests(driverTestConfig);
+
+			runWorkerAuthTests(driverTestConfig);
 		});
 	}
 }

packages/core/src/driver-test-suite/test-inline-client-driver.ts

@@ -43,13 +43,14 @@ export function createTestInlineClientDriver(
 			c: HonoContext | undefined,
 			workerQuery: WorkerQuery,
 			encodingKind: Encoding,
+			params: unknown,
 		): Promise<string> => {
 			return makeInlineRequest<string>(
 				endpoint,
 				encodingKind,
 				transport,
 				"resolveWorkerId",
-				[undefined, workerQuery, encodingKind],
+				[undefined, workerQuery, encodingKind, params],
 			);
 		},