feat: subscribe to events before connect

jog1t · jog1t · commit b1877a0cf600 · 2025-07-29T00:04:40.000+02:00
diff --git a/examples/counter/scripts/connect.ts b/examples/counter/scripts/connect.ts
@@ -8,6 +8,7 @@ async function main() {
 
 	const counter = await client.counter.getOrCreate().connect();
 
+	counter.on("innit", () => console.log("Connected to counter actor"));
 	counter.on("newCount", (count: number) => console.log("Event:", count));
 
 	for (let i = 0; i < 5; i++) {
diff --git a/examples/counter/src/registry.ts b/examples/counter/src/registry.ts
@@ -7,6 +7,9 @@ const counter = actor({
 	onAuth: () => {
 		return true;
 	},
+	onConnect: (c) => {
+		c.broadcast("innit", "Connected to counter actor");
+	},
 	actions: {
 		increment: (c, x: number) => {
 			c.state.count += x;
diff --git a/examples/hono-react/package.json b/examples/hono-react/package.json
@@ -5,7 +5,7 @@
   "type": "module",
   "scripts": {
     "dev": "concurrently \"npm run dev:backend\" \"npm run dev:frontend\"",
-    "dev:backend": "tsx --watch src/backend/server.ts",
+    "dev:backend": "_LOG_LEVEL=1 tsx --watch src/backend/server.ts",
     "dev:frontend": "vite",
     "build": "vite build",
     "check-types": "tsc --noEmit",
diff --git a/examples/hono-react/src/backend/registry.ts b/examples/hono-react/src/backend/registry.ts
@@ -8,6 +8,7 @@ export const counter = actor({
 	actions: {
 		increment: (c, x: number) => {
 			c.state.count += x;
+			console.log("increment", c.state.count);
 			c.broadcast("newCount", c.state.count);
 			return c.state.count;
 		},
diff --git a/packages/core/src/actor/instance.ts b/packages/core/src/actor/instance.ts
@@ -786,6 +786,7 @@ export class ActorInstance<
 		state: CS,
 		driverId: string,
 		driverState: unknown,
+		subs: string[] | undefined,
 		authData: unknown,
 	): Promise<Conn<S, CP, CS, V, I, AD, DB>> {
 		this.#assertReady();
@@ -814,6 +815,12 @@ export class ActorInstance<
 		);
 		this.#connections.set(conn.id, conn);
 
+		if (subs) {
+			for (const sub of subs) {
+				this.#addSubscription(sub, conn, true);
+			}
+		}
+
 		// Add to persistence & save immediately
 		this.#persist.c.push(persist);
 		this.saveState({ immediate: true });
@@ -875,6 +882,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,
@@ -1251,6 +1259,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,
diff --git a/packages/core/src/actor/router-endpoints.ts b/packages/core/src/actor/router-endpoints.ts
@@ -104,6 +104,7 @@ export async function handleWebSocketConnect(
 	actorId: string,
 	encoding: Encoding,
 	parameters: unknown,
+	subs: string[] | undefined,
 	authData: unknown,
 ): Promise<UpgradeWebSocketArgs> {
 	const exposeInternalError = c ? getRequestExposeInternalError(c.req) : false;
@@ -173,6 +174,7 @@ export async function handleWebSocketConnect(
 						connState,
 						CONN_DRIVER_GENERIC_WEBSOCKET,
 						{ encoding } satisfies GenericWebSocketDriverState,
+						subs,
 						authData,
 					);
 
@@ -319,6 +321,7 @@ export async function handleSseConnect(
 	runConfig: RunConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
+	subs: string[] | undefined,
 	authData: unknown,
 ) {
 	const encoding = getRequestEncoding(c.req);
@@ -354,6 +357,7 @@ export async function handleSseConnect(
 				connState,
 				CONN_DRIVER_GENERIC_SSE,
 				{ encoding } satisfies GenericSseDriverState,
+				subs,
 				authData,
 			);
 
@@ -482,6 +486,7 @@ export async function handleAction(
 			connState,
 			CONN_DRIVER_GENERIC_HTTP,
 			{} satisfies GenericHttpDriverState,
+			[],
 			authData,
 		);
 
@@ -698,6 +703,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.
  *
@@ -712,6 +719,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
diff --git a/packages/core/src/actor/router.ts b/packages/core/src/actor/router.ts
@@ -12,6 +12,7 @@ import {
 	HEADER_AUTH_DATA,
 	HEADER_CONN_ID,
 	HEADER_CONN_PARAMS,
+	HEADER_CONN_SUBS,
 	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
 	handleAction,
@@ -83,12 +84,14 @@ 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 ? JSON.parse(subsRaw) : undefined;
 
 				return await handleWebSocketConnect(
 					c as HonoContext,
@@ -97,6 +100,7 @@ export function createActorRouter(
 					c.env.actorId,
 					encoding,
 					connParams,
+					subs,
 					authData,
 				);
 			})(c, noopNext());
@@ -114,8 +118,20 @@ export function createActorRouter(
 		if (authDataRaw) {
 			authData = JSON.parse(authDataRaw);
 		}
+		const subsRaw = c.req.header(HEADER_CONN_SUBS);
+		let subs: string[] | undefined;
+		if (subsRaw) {
+			subs = JSON.parse(subsRaw);
+		}
 
-		return handleSseConnect(c, runConfig, actorDriver, c.env.actorId, authData);
+		return handleSseConnect(
+			c,
+			runConfig,
+			actorDriver,
+			c.env.actorId,
+			subs,
+			authData,
+		);
 	});
 
 	router.post("/action/:action", async (c) => {
diff --git a/packages/core/src/client/actor-conn.ts b/packages/core/src/client/actor-conn.ts
@@ -259,6 +259,7 @@ enc
 			this.#actorQuery,
 			this.#encodingKind,
 			this.#params,
+			["innit"],
 			signal ? { signal } : undefined,
 		);
 		this.#transport = { websocket: ws };
diff --git a/packages/core/src/client/actor-handle.ts b/packages/core/src/client/actor-handle.ts
@@ -48,6 +48,8 @@ export class ActorHandleRaw {
 		this.#params = params;
 	}
 
+	#eventListeners = new Map<string, Set<Function>>();
+
 	/**
 	 * Call a raw action. This method sends an HTTP request to invoke the named action.
 	 *
diff --git a/packages/core/src/client/client.ts b/packages/core/src/client/client.ts
@@ -181,6 +181,7 @@ export interface ClientDriver {
 		actorQuery: ActorQuery,
 		encodingKind: Encoding,
 		params: unknown,
+		subs: string[] | undefined,
 		opts: { signal?: AbortSignal } | undefined,
 	): Promise<WebSocket>;
 	connectSse(
diff --git a/packages/core/src/client/http-client-driver.ts b/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;
 
@@ -133,6 +134,7 @@ export function createHttpClientDriver(managerEndpoint: string): ClientDriver {
 			const protocol = [
 				`query.${encodeURIComponent(JSON.stringify(actorQuery))}`,
 				`encoding.${encodingKind}`,
+				`subs.${encodeURIComponent(JSON.stringify(subs))}`,
 			];
 			if (params)
 				protocol.push(
diff --git a/packages/core/src/drivers/file-system/manager.ts b/packages/core/src/drivers/file-system/manager.ts
@@ -136,6 +136,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
+		subs?: string[],
 	): Promise<WebSocket> {
 		// TODO:
 
@@ -149,6 +150,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				params,
+				subs,
 				undefined,
 			);
 			return new InlineWebSocketAdapter2(wsHandler);
@@ -183,6 +185,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		connParams: unknown,
+		subs: string[] | undefined,
 		authData: unknown,
 	): Promise<Response> {
 		const upgradeWebSocket = this.#runConfig.getUpgradeWebSocket?.();
@@ -198,6 +201,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				connParams,
+				subs,
 				authData,
 			);
 
diff --git a/packages/core/src/manager/driver.ts b/packages/core/src/manager/driver.ts
@@ -33,6 +33,7 @@ export interface ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
+		subs: string[] | undefined,
 		authData: unknown,
 	): Promise<Response>;
 
diff --git a/packages/core/src/manager/protocol/query.ts b/packages/core/src/manager/protocol/query.ts
@@ -64,6 +64,7 @@ export const ConnectWebSocketRequestSchema = z.object({
 	query: ActorQuerySchema.describe("query"),
 	encoding: EncodingSchema.describe("encoding"),
 	connParams: z.unknown().optional().describe("conn_params"),
+	subs: z.array(z.string()).optional().describe("subs"),
 });
 
 export const ConnMessageRequestSchema = z.object({
diff --git a/packages/core/src/manager/router.ts b/packages/core/src/manager/router.ts
@@ -69,10 +69,12 @@ function parseWebSocketProtocols(protocols: string | undefined): {
 	queryRaw: string | undefined;
 	encodingRaw: string | undefined;
 	connParamsRaw: string | undefined;
+	subsRaw: string | undefined;
 } {
 	let queryRaw: string | undefined;
 	let encodingRaw: string | undefined;
 	let connParamsRaw: string | undefined;
+	let subsRaw: string | undefined;
 
 	if (protocols) {
 		const protocolList = protocols.split(",").map((p) => p.trim());
@@ -85,11 +87,15 @@ function parseWebSocketProtocols(protocols: string | undefined): {
 				connParamsRaw = decodeURIComponent(
 					protocol.substring("conn_params.".length),
 				);
+			} else if (protocol.startsWith("subs.")) {
+				// subs are not used in the manager, but we can parse them if needed
+				// this is just for compatibility with the inline client driver
+				subsRaw = decodeURIComponent(protocol.substring("subs.".length));
 			}
 		}
 	}
 
-	return { queryRaw, encodingRaw, connParamsRaw };
+	return { queryRaw, encodingRaw, connParamsRaw, subsRaw };
 }
 
 const OPENAPI_ENCODING = z.string().openapi({
@@ -1211,7 +1217,7 @@ async function handleWebSocketConnectRequest(
 		// Browsers don't support using headers, so this is the only way to
 		// pass data securely.
 		const protocols = c.req.header("sec-websocket-protocol");
-		const { queryRaw, encodingRaw, connParamsRaw } =
+		const { queryRaw, encodingRaw, connParamsRaw, subsRaw } =
 			parseWebSocketProtocols(protocols);
 
 		// Parse query
@@ -1236,13 +1242,30 @@ async function handleWebSocketConnectRequest(
 			);
 		}
 
+		// Parse subscriptions
+		let subs: string[] | undefined;
+		if (subsRaw) {
+			try {
+				subs = JSON.parse(subsRaw!);
+				if (!Array.isArray(subs)) {
+					throw new Error("Subscriptions must be an array");
+				}
+			} catch (error) {
+				logger().error("invalid subscriptions json", { error });
+				throw new errors.InvalidRequest(
+					`Invalid subscriptions JSON: ${stringifyError(error)}`,
+				);
+			}
+		}
+
 		// We can't use the standard headers with WebSockets
 		//
 		// All other information will be sent over the socket itself, since that data needs to be E2EE
 		const params = ConnectWebSocketRequestSchema.safeParse({
 			query: queryUnvalidated,
 			encoding: encodingRaw,
 			connParams: connParamsUnvalidated,
+			subs,
 		});
 		if (!params.success) {
 			logger().error("invalid connection parameters", {
@@ -1281,6 +1304,7 @@ async function handleWebSocketConnectRequest(
 			actorId,
 			params.data.encoding,
 			params.data.connParams,
+			params.data.subs,
 			authData,
 		);
 	} catch (error) {
@@ -1623,6 +1647,7 @@ async function handleRawWebSocketRequest(
 		const {
 			queryRaw: queryFromProtocol,
 			connParamsRaw: connParamsFromProtocol,
+			subsRaw: subsFromProtocol,
 		} = parseWebSocketProtocols(protocols);
 
 		if (!queryFromProtocol) {
@@ -1636,6 +1661,11 @@ async function handleRawWebSocketRequest(
 			connParams = JSON.parse(connParamsFromProtocol);
 		}
 
+		let subs: string[] | undefined;
+		if (subsFromProtocol) {
+			subs = JSON.parse(subsFromProtocol);
+		}
+
 		// Authenticate the request
 		const authData = await authenticateEndpoint(
 			c,
@@ -1671,6 +1701,7 @@ async function handleRawWebSocketRequest(
 			actorId,
 			"json", // Default encoding for raw WebSocket
 			connParams,
+			subs,
 			authData,
 		);
 	} catch (error) {

Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,8 @@ export class ActorHandleRaw {`
`48`	`48`	`this.#params = params;`
`49`	`49`	`}`
`50`	`50`
	`51`	`+ #eventListeners = new Map<string, Set<Function>>();`
	`52`	`+`
`51`	`53`	`/**`
`52`	`54`	`* Call a raw action. This method sends an HTTP request to invoke the named action.`
`53`	`55`	`*`