refactor(rivetkit): restructure connection lifecycle contexts

Author: NathanFlurry

examples/cursors-raw-websocket/src/backend/registry.ts

@@ -2,7 +2,7 @@ import { actor } from "rivetkit";
 
 export const counterWithParams = actor({
 	state: { count: 0, initializers: [] as string[] },
-	createConnState: (c, opts, params: { name?: string }) => {
+	createConnState: (c, params: { name?: string }) => {
 		return {
 			name: params.name || "anonymous",
 		};

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-params.ts

@@ -16,7 +16,6 @@ export const connStateActor = actor({
 	// Define connection state
 	createConnState: (
 		c,
-		opts,
 		params: { username?: string; role?: string; noCount?: boolean },
 	): ConnState => {
 		return {

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-state.ts

@@ -7,13 +7,13 @@ export const counterWithLifecycle = actor({
 		count: 0,
 		events: [] as string[],
 	},
-	createConnState: (c, opts, params: ConnParams) => ({
+	createConnState: (c, params: ConnParams) => ({
 		joinTime: Date.now(),
 	}),
 	onWake: (c) => {
 		c.state.events.push("onWake");
 	},
-	onBeforeConnect: (c, opts, params: ConnParams) => {
+	onBeforeConnect: (c, params: ConnParams) => {
 		if (params?.trackLifecycle) c.state.events.push("onBeforeConnect");
 	},
 	onConnect: (c, conn) => {

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/lifecycle.ts

@@ -5,7 +5,7 @@ export const rawWebSocketActor = actor({
 		connectionCount: 0,
 		messageCount: 0,
 	},
-	onWebSocket(ctx, websocket, opts) {
+	onWebSocket(ctx, websocket) {
 		ctx.state.connectionCount = ctx.state.connectionCount + 1;
 		console.log(
 			`[ACTOR] New connection, count: ${ctx.state.connectionCount}`,
@@ -51,15 +51,16 @@ export const rawWebSocketActor = actor({
 							}),
 						);
 					} else if (parsed.type === "getRequestInfo") {
+						throw "TODO";
 						// Send back the request URL info
-						websocket.send(
-							JSON.stringify({
-								type: "requestInfo",
-								url: opts.request.url,
-								pathname: new URL(opts.request.url).pathname,
-								search: new URL(opts.request.url).search,
-							}),
-						);
+						// websocket.send(
+						// 	JSON.stringify({
+						// 		type: "requestInfo",
+						// 		url: opts.request.url,
+						// 		pathname: new URL(opts.request.url).pathname,
+						// 		search: new URL(opts.request.url).search,
+						// 	}),
+						// );
 					} else {
 						// Echo back
 						websocket.send(data);
@@ -93,7 +94,7 @@ export const rawWebSocketActor = actor({
 });
 
 export const rawWebSocketBinaryActor = actor({
-	onWebSocket(ctx, websocket, opts) {
+	onWebSocket(ctx, websocket) {
 		// Handle binary data
 		websocket.addEventListener("message", (event: any) => {
 			const data = event.data;

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts

@@ -31,19 +31,19 @@ export const requestAccessActor = actor({
 			requestHeaders: {} as Record<string, string>,
 		},
 	},
-	createConnState: (c, { request }, params: { trackRequest?: boolean }) => {
+	createConnState: (c, params: { trackRequest?: boolean }) => {
 		// In createConnState, the state isn't available yet.
 
 		return {
 			trackRequest: params?.trackRequest || false,
 			requestInfo:
-				params?.trackRequest && request
+				params?.trackRequest && c.request
 					? {
 							hasRequest: true,
-							requestUrl: request.url,
-							requestMethod: request.method,
+							requestUrl: c.request.url,
+							requestMethod: c.request.method,
 							requestHeaders: Object.fromEntries(
-								request.headers.entries(),
+								c.request.headers.entries(),
 							),
 						}
 					: null,
@@ -55,16 +55,16 @@ export const requestAccessActor = actor({
 			c.state.createConnStateRequest = conn.state.requestInfo;
 		}
 	},
-	onBeforeConnect: (c, { request }, params) => {
+	onBeforeConnect: (c, params) => {
 		if (params?.trackRequest) {
-			if (request) {
+			if (c.request) {
 				c.state.onBeforeConnectRequest.hasRequest = true;
-				c.state.onBeforeConnectRequest.requestUrl = request.url;
-				c.state.onBeforeConnectRequest.requestMethod = request.method;
+				c.state.onBeforeConnectRequest.requestUrl = c.request.url;
+				c.state.onBeforeConnectRequest.requestMethod = c.request.method;
 
 				// Store select headers
 				const headers: Record<string, string> = {};
-				request.headers.forEach((value, key) => {
+				c.request.headers.forEach((value, key) => {
 					headers[key] = value;
 				});
 				c.state.onBeforeConnectRequest.requestHeaders = headers;
@@ -101,15 +101,16 @@ export const requestAccessActor = actor({
 			},
 		);
 	},
-	onWebSocket: (c, websocket, { request }) => {
+	onWebSocket: (c, websocket) => {
+		if (!c.request) throw "Missing request";
 		// Store request info
 		c.state.onWebSocketRequest.hasRequest = true;
-		c.state.onWebSocketRequest.requestUrl = request.url;
-		c.state.onWebSocketRequest.requestMethod = request.method;
+		c.state.onWebSocketRequest.requestUrl = c.request.url;
+		c.state.onWebSocketRequest.requestMethod = c.request.method;
 
 		// Store select headers
 		const headers: Record<string, string> = {};
-		request.headers.forEach((value, key) => {
+		c.request.headers.forEach((value, key) => {
 			headers[key] = value;
 		});
 		c.state.onWebSocketRequest.requestHeaders = headers;
@@ -118,8 +119,8 @@ export const requestAccessActor = actor({
 		websocket.send(
 			JSON.stringify({
 				hasRequest: true,
-				requestUrl: request.url,
-				requestMethod: request.method,
+				requestUrl: c.request.url,
+				requestMethod: c.request.method,
 				requestHeaders: headers,
 			}),
 		);

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/request-access.ts

@@ -112,7 +112,7 @@ export const sleepWithRawWebSocket = actor({
 	onSleep: (c) => {
 		c.state.sleepCount += 1;
 	},
-	onWebSocket: (c, websocket: UniversalWebSocket, opts) => {
+	onWebSocket: (c, websocket: UniversalWebSocket) => {
 		c.state.connectionCount += 1;
 		c.log.info({
 			msg: "websocket connected",

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep.ts

@@ -3,6 +3,9 @@ import type { UniversalWebSocket } from "@/common/websocket-interface";
 import type { Conn } from "./conn/mod";
 import type { ActionContext } from "./contexts/action";
 import type { ActorContext } from "./contexts/actor";
+import type { CreateConnStateContext } from "./contexts/create-conn-state";
+import type { OnBeforeConnectContext } from "./contexts/on-before-connect";
+import type { OnConnectContext } from "./contexts/on-connect";
 import type { RequestContext } from "./contexts/request";
 import type { WebSocketContext } from "./contexts/websocket";
 import type { AnyDatabaseProvider } from "./database";
@@ -113,15 +116,6 @@ export const ActorConfigSchema = z
 		},
 	);
 
-export interface OnConnectOptions {
-	/**
-	 * The request object associated with the connection.
-	 *
-	 * @experimental
-	 */
-	request?: Request;
-}
-
 // Creates state config
 //
 // This must have only one or the other or else TState will not be able to be inferred
@@ -148,13 +142,12 @@ type CreateConnState<
 	TConnState,
 	TVars,
 	TInput,
-	TDatabase,
+	TDatabase extends AnyDatabaseProvider,
 > =
 	| { connState: TConnState }
 	| {
 			createConnState: (
-				c: InitContext,
-				opts: OnConnectOptions,
+				c: CreateConnStateContext<TState, TVars, TInput, TDatabase>,
 				params: TConnParams,
 			) => TConnState | Promise<TConnState>;
 	  }
@@ -323,15 +316,7 @@ interface BaseActorConfig<
 	 * @throws Throw an error to reject the connection
 	 */
 	onBeforeConnect?: (
-		c: ActorContext<
-			TState,
-			TConnParams,
-			TConnState,
-			TVars,
-			TInput,
-			TDatabase
-		>,
-		opts: OnConnectOptions,
+		c: OnBeforeConnectContext<TState, TVars, TInput, TDatabase>,
 		params: TConnParams,
 	) => void | Promise<void>;
 
@@ -345,14 +330,7 @@ interface BaseActorConfig<
 	 * @returns Void or a Promise that resolves when connection handling is complete
 	 */
 	onConnect?: (
-		c: ActorContext<
-			TState,
-			TConnParams,
-			TConnState,
-			TVars,
-			TInput,
-			TDatabase
-		>,
+		c: OnConnectContext<TState, TVars, TInput, TDatabase>,
 		conn: Conn<TState, TConnParams, TConnState, TVars, TInput, TDatabase>,
 	) => void | Promise<void>;
 
@@ -446,7 +424,6 @@ interface BaseActorConfig<
 			TDatabase
 		>,
 		websocket: UniversalWebSocket,
-		opts: { request: Request },
 	) => void | Promise<void>;
 
 	actions: TActions;

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

@@ -5,7 +5,7 @@ import type { CachedSerializer, Encoding } from "@/actor/protocol/serde";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import { type ConnDriver, DriverReadyState } from "../driver";
 
-export type ConnDriverWebSocketState = {};
+export type ConnDriverWebSocketState = Record<never, never>;
 
 export function createWebSocketSocket(
 	requestId: string,