From f8d641c8b40b772b457687dd3716a3e4784881a9 Mon Sep 17 00:00:00 2001
From: Kacper Wojciechowski <39823706+jog1t@users.noreply.github.com>
Date: Mon, 28 Jul 2025 20:52:22 +0200
Subject: [PATCH] fix(react): memoize event listener

---
 packages/frameworks/framework-base/lib/mod.ts | 36 ++++++++++---------
 packages/frameworks/react/src/mod.tsx         | 22 ++++++++----
 2 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/packages/frameworks/framework-base/lib/mod.ts b/packages/frameworks/framework-base/lib/mod.ts
index 6eb150edd..0cbc16add 100644
--- a/packages/frameworks/framework-base/lib/mod.ts
+++ b/packages/frameworks/framework-base/lib/mod.ts
@@ -7,7 +7,6 @@ import type {
 } from "@rivetkit/core/client";
 import { Derived, Effect, Store, type Updater } from "@tanstack/store";
 
-// biome-ignore lint/suspicious/noExplicitAny: its a generic actor registry
 export type AnyActorRegistry = Registry<any>;
 
 interface ActorStateReference<AD extends AnyActorDefinition> {
@@ -113,11 +112,27 @@ export interface ActorOptions<
 	enabled?: boolean;
 }
 
-// biome-ignore lint/suspicious/noExplicitAny: actor name can be anything
+export type ActorsStateDerived<
+	Registry extends AnyActorRegistry,
+	WorkerName extends keyof ExtractActorsFromRegistry<Registry>,
+> = Derived<
+	Omit<
+		InternalRivetKitStore<
+			Registry,
+			ExtractActorsFromRegistry<Registry>
+		>["actors"][string],
+		"handle" | "connection"
+	> & {
+		handle: ActorHandle<ExtractActorsFromRegistry<Registry>[WorkerName]> | null;
+		connection: ActorConn<
+			ExtractActorsFromRegistry<Registry>[WorkerName]
+		> | null;
+	}
+>;
+
 export type AnyActorOptions = ActorOptions<AnyActorRegistry, any>;
 
 export interface CreateRivetKitOptions<Registry extends AnyActorRegistry> {
-	// biome-ignore lint/suspicious/noExplicitAny: actor name can be anything
 	hashFunction?: (opts: ActorOptions<Registry, any>) => string;
 }
 
@@ -144,7 +159,6 @@ export function createRivetKit<
 			create: () => void;
 			addEventListener?: (
 				event: string,
-				// biome-ignore lint/suspicious/noExplicitAny: need any specific type here
 				handler: (...args: any[]) => void,
 			) => void;
 		}
@@ -158,12 +172,7 @@ export function createRivetKit<
 		if (cached) {
 			return {
 				...cached,
-				state: cached.state as Derived<
-					Omit<RivetKitStore["actors"][string], "handle" | "connection"> & {
-						handle: ActorHandle<Actors[ActorName]> | null;
-						connection: ActorConn<Actors[ActorName]> | null;
-					}
-				>,
+				state: cached.state as ActorsStateDerived<Registry, ActorName>,
 			};
 		}
 
@@ -328,12 +337,7 @@ export function createRivetKit<
 		return {
 			mount,
 			setState,
-			state: derived as Derived<
-				Omit<RivetKitStore["actors"][string], "handle" | "connection"> & {
-					handle: ActorHandle<Actors[ActorName]> | null;
-					connection: ActorConn<Actors[ActorName]> | null;
-				}
-			>,
+			state: derived as ActorsStateDerived<Registry, ActorName>,
 			create,
 			key,
 		};
diff --git a/packages/frameworks/react/src/mod.tsx b/packages/frameworks/react/src/mod.tsx
index 81c8c81ad..7ffce4672 100644
--- a/packages/frameworks/react/src/mod.tsx
+++ b/packages/frameworks/react/src/mod.tsx
@@ -5,8 +5,8 @@ import {
 	type CreateRivetKitOptions,
 	createRivetKit as createVanillaRivetKit,
 } from "@rivetkit/framework-base";
+import { useEffect, useRef } from "react";
 import { useStore } from "@tanstack/react-store";
-import { useEffect } from "react";
 
 export { createClient } from "@rivetkit/core/client";
 
@@ -57,19 +57,29 @@ export function createRivetKit<Registry extends AnyActorRegistry>(
 		 * @param eventName The name of the event to listen for.
 		 * @param handler The function to call when the event is emitted.
 		 */
-		const useEvent = (
+		function useEvent(
 			eventName: string,
 			// biome-ignore lint/suspicious/noExplicitAny: strong typing of handler is not supported yet
 			handler: (...args: any[]) => void,
-		) => {
+		) {
+			const ref = useRef(handler);
+			const actorState = useStore(state) || {};
+
+			useEffect(() => {
+				ref.current = handler;
+			}, [handler]);
+
 			// biome-ignore lint/correctness/useExhaustiveDependencies: it's okay to not include all dependencies here
 			useEffect(() => {
 				if (!actorState?.connection) return;
 
-				const connection = actorState.connection;
-				return connection.on(eventName, handler);
-			}, [actorState.connection, actorState.isConnected, eventName, handler]);
+				function eventHandler(...args: any[]) {
+					ref.current(...args);
+				}
+				return actorState.connection.on(eventName, eventHandler);
+			}, [actorState.connection, actorState.isConnected, actorState.hash, eventName]);
 		};
+
 		return {
 			...actorState,
 			useEvent,