feat: subscribe to events before connect

Author: jog1t

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

@@ -3,6 +3,10 @@ import { actor } from "@rivetkit/core";
 export const counter = actor({
 	onAuth: () => {},
 	state: { count: 0 },
+	onConnect: (c, conn) => {
+		c.broadcast("onconnect:broadcast", "Hello!");
+		conn.send("onconnect:msg", "Welcome to the counter actor!");
+	},
 	actions: {
 		increment: (c, x: number) => {
 			c.state.count += x;

packages/core/src/actor/instance.ts

@@ -811,6 +811,7 @@ export class ActorInstance<
 		state: CS,
 		driverId: ConnectionDriver,
 		driverState: unknown,
+		subscriptions: string[],
 		authData: unknown,
 	): Promise<Conn<S, CP, CS, V, I, AD, DB>> {
 		this.#assertReady();
@@ -840,6 +841,12 @@ export class ActorInstance<
 		);
 		this.#connections.set(conn.id, conn);
 
+		if (subscriptions) {
+			for (const sub of subscriptions) {
+				this.#addSubscription(sub, conn, true);
+			}
+		}
+
 		// Add to persistence & save immediately
 		this.#persist.c.push(persist);
 		this.saveState({ immediate: true });
@@ -901,6 +908,7 @@ export class ActorInstance<
 				return await this.executeAction(ctx, name, args);
 			},
 			onSubscribe: async (eventName, conn) => {
+				console.log("subscribing to event", { eventName, connId: conn.id });
 				this.inspector.emitter.emit("eventFired", {
 					type: "subscribe",
 					eventName,
@@ -1334,6 +1342,13 @@ export class ActorInstance<
 	_broadcast<Args extends Array<unknown>>(name: string, ...args: Args) {
 		this.#assertReady();
 
+		console.log("broadcasting event", {
+			name,
+			args,
+			actorId: this.id,
+			subscriptions: this.#subscriptionIndex.size,
+			connections: this.conns.size,
+		});
 		this.inspector.emitter.emit("eventFired", {
 			type: "broadcast",
 			eventName: name,

packages/core/src/actor/protocol/serde.ts

@@ -17,6 +17,9 @@ export const EncodingSchema = z.enum(["json", "cbor"]);
  */
 export type Encoding = z.infer<typeof EncodingSchema>;
 
+export const SubscriptionsListSchema = z.array(z.string());
+export type SubscriptionsList = z.infer<typeof SubscriptionsListSchema>;
+
 /**
  * Helper class that helps serialize data without re-serializing for the same encoding.
  */

packages/core/src/actor/router-endpoints.ts

@@ -108,6 +108,7 @@ export async function handleWebSocketConnect(
 	encoding: Encoding,
 	parameters: unknown,
 	authData: unknown,
+	subscriptions: string[],
 ): Promise<UpgradeWebSocketArgs> {
 	const exposeInternalError = c ? getRequestExposeInternalError(c.req) : false;
 
@@ -176,6 +177,7 @@ export async function handleWebSocketConnect(
 						connState,
 						CONNECTION_DRIVER_WEBSOCKET,
 						{ encoding } satisfies GenericWebSocketDriverState,
+						subscriptions,
 						authData,
 					);
 
@@ -322,6 +324,7 @@ export async function handleSseConnect(
 	runConfig: RunConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
+	subscriptions: string[],
 	authData: unknown,
 ) {
 	const encoding = getRequestEncoding(c.req);
@@ -357,6 +360,7 @@ export async function handleSseConnect(
 				connState,
 				CONNECTION_DRIVER_SSE,
 				{ encoding } satisfies GenericSseDriverState,
+				subscriptions,
 				authData,
 			);
 
@@ -487,6 +491,7 @@ export async function handleAction(
 			connState,
 			CONNECTION_DRIVER_HTTP,
 			{} satisfies GenericHttpDriverState,
+			[],
 			authData,
 		);
 
@@ -706,6 +711,8 @@ export const HEADER_CONN_ID = "X-RivetKit-Conn";
 
 export const HEADER_CONN_TOKEN = "X-RivetKit-Conn-Token";
 
+export const HEADER_CONN_SUBS = "X-RivetKit-Conn-Subs";
+
 /**
  * Headers that publics can send from public clients.
  *
@@ -720,6 +727,7 @@ export const ALLOWED_PUBLIC_HEADERS = [
 	HEADER_ACTOR_ID,
 	HEADER_CONN_ID,
 	HEADER_CONN_TOKEN,
+	HEADER_CONN_SUBS,
 ];
 
 // Helper to get connection parameters for the request

packages/core/src/actor/router.ts

@@ -1,6 +1,9 @@
 import { Hono, type Context as HonoContext } from "hono";
 import invariant from "invariant";
-import { EncodingSchema } from "@/actor/protocol/serde";
+import {
+	EncodingSchema,
+	SubscriptionsListSchema,
+} from "@/actor/protocol/serde";
 import {
 	type ActionOpts,
 	type ActionOutput,
@@ -12,6 +15,7 @@ import {
 	HEADER_AUTH_DATA,
 	HEADER_CONN_ID,
 	HEADER_CONN_PARAMS,
+	HEADER_CONN_SUBS,
 	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
 	handleAction,
@@ -83,12 +87,16 @@ export function createActorRouter(
 				const encodingRaw = c.req.header(HEADER_ENCODING);
 				const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
 				const authDataRaw = c.req.header(HEADER_AUTH_DATA);
+				const subsRaw = c.req.header(HEADER_CONN_SUBS);
 
 				const encoding = EncodingSchema.parse(encodingRaw);
 				const connParams = connParamsRaw
 					? JSON.parse(connParamsRaw)
 					: undefined;
 				const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
+				const subs = subsRaw
+					? SubscriptionsListSchema.parse(JSON.parse(subsRaw))
+					: [];
 
 				return await handleWebSocketConnect(
 					c as HonoContext,
@@ -97,6 +105,7 @@ export function createActorRouter(
 					c.env.actorId,
 					encoding,
 					connParams,
+					subs,
 					authData,
 				);
 			})(c, noopNext());
@@ -114,8 +123,20 @@ export function createActorRouter(
 		if (authDataRaw) {
 			authData = JSON.parse(authDataRaw);
 		}
+		const subsRaw = c.req.header(HEADER_CONN_SUBS);
+
+		const subscriptions = subsRaw
+			? SubscriptionsListSchema.parse(JSON.parse(subsRaw))
+			: [];
 
-		return handleSseConnect(c, runConfig, actorDriver, c.env.actorId, authData);
+		return handleSseConnect(
+			c,
+			runConfig,
+			actorDriver,
+			c.env.actorId,
+			subscriptions,
+			authData,
+		);
 	});
 
 	router.post("/action/:action", async (c) => {

packages/core/src/client/actor-conn.ts

@@ -98,7 +98,7 @@ export class ActorConnRaw {
 	/**
 	 * Interval that keeps the NodeJS process alive if this is the only thing running.
 	 *
-	 * See ttps://github.com/nodejs/node/issues/22088
+	 * @see https://github.com/nodejs/node/issues/22088
 	 */
 	#keepNodeAliveInterval: NodeJS.Timeout;
 
@@ -111,8 +111,6 @@ export class ActorConnRaw {
 	#encodingKind: Encoding;
 	#actorQuery: ActorQuery;
 
-	// TODO: ws message queue
-
 	/**
 	 * Do not call this directly.
 	 *
@@ -187,7 +185,6 @@ export class ActorConnRaw {
 
 	/**
 	 * Do not call this directly.
-enc
 	 * Establishes a connection to the server using the specified endpoint & encoding & driver.
 	 *
 	 * @protected
@@ -259,6 +256,7 @@ enc
 			this.#actorQuery,
 			this.#encodingKind,
 			this.#params,
+			Array.from(this.#eventSubscriptions.keys()),
 			signal ? { signal } : undefined,
 		);
 		this.#transport = { websocket: ws };
@@ -282,6 +280,7 @@ enc
 			this.#actorQuery,
 			this.#encodingKind,
 			this.#params,
+			Array.from(this.#eventSubscriptions.keys()),
 			signal ? { signal } : undefined,
 		);
 		this.#transport = { sse: eventSource };
@@ -807,3 +806,15 @@ enc
  */
 export type ActorConn<AD extends AnyActorDefinition> = ActorConnRaw &
 	ActorDefinitionActions<AD>;
+
+/**
+ * Connection to a actor. Allows calling actor's remote procedure calls with inferred types. See {@link ActorConnRaw} for underlying methods.
+ * Needs to be established manually using #connect.
+ *
+ * @template AD The actor class that this connection is for.
+ * @see {@link ActorConnRaw}
+ * @see {@link ActorConn}
+ */
+export type ActorManualConn<AD extends AnyActorDefinition> = ActorConnRaw & {
+	connect: () => void;
+} & ActorDefinitionActions<AD>;

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

@@ -5,11 +5,16 @@ import { assertUnreachable } from "@/actor/utils";
 import { importWebSocket } from "@/common/websocket";
 import type { ActorQuery } from "@/manager/protocol/query";
 import type { ActorDefinitionActions } from "./actor-common";
-import { type ActorConn, ActorConnRaw } from "./actor-conn";
+import {
+	type ActorConn,
+	ActorConnRaw,
+	type ActorManualConn,
+} from "./actor-conn";
 import {
 	type ClientDriver,
 	type ClientRaw,
 	CREATE_ACTOR_CONN_PROXY,
+	CREATE_ACTOR_PROXY,
 } from "./client";
 import { logger } from "./log";
 import { rawHttpFetch, rawWebSocket } from "./raw-utils";
@@ -98,6 +103,32 @@ export class ActorHandleRaw {
 		) as ActorConn<AnyActorDefinition>;
 	}
 
+	/**
+	 * Creates a new connection to the actor, that should be manually connected.
+	 * This is useful for creating connections that are not immediately connected,
+	 * such as when you want to set up event listeners before connecting.
+	 *
+	 * @param AD - The actor definition for the connection.
+	 * @returns {ActorConn<AD>} A connection to the actor.
+	 */
+	create(): ActorManualConn<AnyActorDefinition> {
+		logger().debug("creating a connection from handle", {
+			query: this.#actorQuery,
+		});
+
+		const conn = new ActorConnRaw(
+			this.#client,
+			this.#driver,
+			this.#params,
+			this.#encodingKind,
+			this.#actorQuery,
+		);
+
+		return this.#client[CREATE_ACTOR_PROXY](
+			conn,
+		) as ActorManualConn<AnyActorDefinition>;
+	}
+
 	/**
 	 * Makes a raw HTTP request to the actor.
 	 *
@@ -188,10 +219,12 @@ export class ActorHandleRaw {
  */
 export type ActorHandle<AD extends AnyActorDefinition> = Omit<
 	ActorHandleRaw,
-	"connect"
+	"connect" | "create"
 > & {
 	// Add typed version of ActorConn (instead of using AnyActorDefinition)
 	connect(): ActorConn<AD>;
 	// Resolve method returns the actor ID
 	resolve(): Promise<string>;
+	// Add typed version of create
+	create(): ActorManualConn<AD>;
 } & ActorDefinitionActions<AD>;

packages/core/src/client/client.ts

@@ -11,6 +11,7 @@ import type { ActorActionFunction } from "./actor-common";
 import {
 	type ActorConn,
 	type ActorConnRaw,
+	type ActorManualConn,
 	CONNECT_SYMBOL,
 } from "./actor-conn";
 import { type ActorHandle, ActorHandleRaw } from "./actor-handle";
@@ -157,6 +158,7 @@ export interface Region {
 
 export const ACTOR_CONNS_SYMBOL = Symbol("actorConns");
 export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy");
+export const CREATE_ACTOR_PROXY = Symbol("createActorProxy");
 export const TRANSPORT_SYMBOL = Symbol("transport");
 
 export interface ClientDriver {
@@ -181,13 +183,15 @@ export interface ClientDriver {
 		actorQuery: ActorQuery,
 		encodingKind: Encoding,
 		params: unknown,
+		subscriptions: string[],
 		opts: { signal?: AbortSignal } | undefined,
 	): Promise<WebSocket>;
 	connectSse(
 		c: HonoContext | undefined,
 		actorQuery: ActorQuery,
 		encodingKind: Encoding,
 		params: unknown,
+		subscriptions: string[],
 		opts: { signal?: AbortSignal } | undefined,
 	): Promise<UniversalEventSource>;
 	sendHttpMessage(
@@ -426,12 +430,35 @@ export class ClientRaw {
 		// Save to connection list
 		this[ACTOR_CONNS_SYMBOL].add(conn);
 
+		logger().debug("creating actor proxy for connection and connecting", {
+			conn,
+		});
+
 		// Start connection
 		conn[CONNECT_SYMBOL]();
 
 		return createActorProxy(conn) as ActorConn<AD>;
 	}
 
+	[CREATE_ACTOR_PROXY]<AD extends AnyActorDefinition>(
+		conn: ActorConnRaw,
+	): ActorConn<AD> {
+		// Save to connection list
+		this[ACTOR_CONNS_SYMBOL].add(conn);
+
+		logger().debug("creating actor proxy for connection", {
+			conn,
+		});
+
+		Object.assign(conn, {
+			connect: () => {
+				conn[CONNECT_SYMBOL]();
+			},
+		});
+
+		return createActorProxy(conn) as ActorManualConn<AD>;
+	}
+
 	/**
 	 * Disconnects from all actors.
 	 *

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

@@ -121,6 +121,7 @@ export function createHttpClientDriver(managerEndpoint: string): ClientDriver {
 			actorQuery: ActorQuery,
 			encodingKind: Encoding,
 			params: unknown,
+			subs: string[] | undefined,
 		): Promise<WebSocket> => {
 			const { WebSocket } = await dynamicImports;
 
@@ -138,6 +139,9 @@ export function createHttpClientDriver(managerEndpoint: string): ClientDriver {
 				protocol.push(
 					`conn_params.${encodeURIComponent(JSON.stringify(params))}`,
 				);
+			if (subs) {
+				protocol.push(`subs.${encodeURIComponent(JSON.stringify(subs))}`);
+			}
 
 			// HACK: See packages/drivers/cloudflare-workers/src/websocket.ts
 			protocol.push("rivetkit");