fix(rivetkit): unwrap state proxies before persisting to stop exponential proxy nesting

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

@@ -1,3 +1,4 @@
+import onChange from "@rivetkit/on-change";
 import { actor } from "rivetkit";
 
 export type ConnState = {
@@ -6,12 +7,33 @@ export type ConnState = {
 	counter: number;
 	createdAt: number;
 	noCount: boolean;
+	capabilities: { tags: string[] };
 };
 
+/**
+ * Counts how many write-through proxy layers wrap a value. A value read off a
+ * state proxy is wrapped exactly once; more layers mean previously read
+ * proxies were persisted back into state.
+ */
+function proxyDepth(value: unknown): number {
+	let depth = 0;
+	let current = value;
+	while (current !== null && typeof current === "object") {
+		const target = onChange.target(current as Record<string, any>);
+		if (target === current) {
+			break;
+		}
+		depth++;
+		current = target;
+	}
+	return depth;
+}
+
 export const connStateActor = actor({
 	state: {
 		sharedCounter: 0,
 		disconnectionCount: 0,
+		nested: { tags: ["read", "write"] as string[] },
 	},
 	// Define connection state
 	createConnState: (
@@ -24,6 +46,7 @@ export const connStateActor = actor({
 			counter: 0,
 			createdAt: Date.now(),
 			noCount: params?.noCount ?? false,
+			capabilities: { tags: ["read", "write"] },
 		};
 	},
 	// Lifecycle hook when a connection is established
@@ -118,6 +141,29 @@ export const connStateActor = actor({
 			if (updates.role) c.conn.state.role = updates.role;
 			return c.conn.state;
 		},
+		// Replacing state with a spread of the current state is the common
+		// update pattern. Each read hands back a deep write-through proxy, so
+		// the nested values in the spread are proxies themselves.
+		spreadUpdateConnState: (c, iterations: number) => {
+			for (let i = 0; i < iterations; i++) {
+				c.conn.state = { ...c.conn.state, counter: i };
+			}
+			return {
+				depth: proxyDepth(c.conn.state.capabilities),
+				tags: [...c.conn.state.capabilities.tags],
+			};
+		},
+
+		spreadUpdateActorState: (c, iterations: number) => {
+			for (let i = 0; i < iterations; i++) {
+				c.state = { ...c.state, sharedCounter: i };
+			}
+			return {
+				depth: proxyDepth(c.state.nested),
+				tags: [...c.state.nested.tags],
+			};
+		},
+
 		disconnectSelf: (c, reason?: string) => {
 			c.conn.disconnect(reason ?? "test.disconnect");
 			return true;

rivetkit-typescript/packages/rivetkit/src/registry/native.ts

@@ -104,7 +104,10 @@ import type {
 	WebSocketHandle,
 } from "./runtime";
 import { loadWasmRuntime } from "./wasm-runtime";
-import { createWriteThroughProxy } from "./write-through-proxy";
+import {
+	createWriteThroughProxy,
+	unwrapWriteThroughProxy,
+} from "./write-through-proxy";
 
 const textEncoder = new TextEncoder();
 const textDecoder = new TextDecoder();
@@ -126,6 +129,10 @@ type NativeOnStateChangeHandler = (
 ) => void | Promise<void>;
 type NativePersistConnState = {
 	state: unknown;
+	// Memoized deep write-through proxy and the state object it wraps. Rebuilt
+	// only when the underlying state object identity changes.
+	stateProxy?: unknown;
+	stateProxyTarget?: unknown;
 };
 
 const defaultRuntimeLoaders: RuntimeLoaders = {
@@ -1270,8 +1277,32 @@ class NativeConnAdapter {
 
 	get state(): unknown {
 		const nextState = this.#readState();
+		if (!this.#ctx) {
+			return this.#createStateProxy(nextState);
+		}
+
+		// Reading `conn.state` rebuilds the deep write-through proxy, which
+		// allocates fresh on-change caches and rewraps the whole tree. Memoize
+		// the proxy keyed on the underlying state object so repeated reads and
+		// deep read cascades reuse a single proxy.
+		const connState = getNativeConnPersistState(
+			this.#runtime,
+			this.#ctx,
+			this.#conn,
+		);
+		if (
+			connState.stateProxy === undefined ||
+			connState.stateProxyTarget !== nextState
+		) {
+			connState.stateProxyTarget = nextState;
+			connState.stateProxy = this.#createStateProxy(nextState);
+		}
+		return connState.stateProxy;
+	}
+
+	#createStateProxy(state: unknown): unknown {
 		return createWriteThroughProxy(
-			nextState,
+			state,
 			(nextValue) => {
 				this.#writeState(nextValue, { writeNative: true });
 			},
@@ -1282,12 +1313,15 @@ class NativeConnAdapter {
 	}
 
 	set state(value: unknown) {
-		assertJsonCompatValue(value);
-		this.#writeState(value, { writeNative: true });
+		const nextValue = unwrapWriteThroughProxy(value);
+		assertJsonCompatValue(nextValue);
+		this.#writeState(nextValue, { writeNative: true });
 	}
 
 	initializeState(value: unknown): void {
-		this.#writeState(value, { writeNative: false });
+		this.#writeState(unwrapWriteThroughProxy(value), {
+			writeNative: false,
+		});
 	}
 
 	get isHibernatable(): boolean {
@@ -2693,15 +2727,18 @@ export class ActorContextHandleAdapter {
 			throw stateNotEnabledError();
 		}
 		this.#assertCanMutateState();
-		assertJsonCompatValue(value);
-		this.#writeState(value, { scheduleSave: true });
+		const nextValue = unwrapWriteThroughProxy(value);
+		assertJsonCompatValue(nextValue);
+		this.#writeState(nextValue, { scheduleSave: true });
 	}
 
 	initializeState(value: unknown): void {
 		if (!this.#stateEnabled) {
 			return;
 		}
-		this.#writeState(value, { scheduleSave: false });
+		this.#writeState(unwrapWriteThroughProxy(value), {
+			scheduleSave: false,
+		});
 	}
 
 	get vars(): unknown {

rivetkit-typescript/packages/rivetkit/src/registry/write-through-proxy.ts

@@ -38,3 +38,104 @@ export function createWriteThroughProxy<T>(
 		},
 	) as T;
 }
+
+/**
+ * Returns the raw target behind an `@rivetkit/on-change` proxy, following
+ * chains of proxies wrapping proxies until reaching a plain value.
+ */
+function unwrapProxy(value: unknown): unknown {
+	let current = value;
+	while (current !== null && typeof current === "object") {
+		const target = onChange.target(current as Record<string, any>);
+		if (target === current) {
+			break;
+		}
+		current = target;
+	}
+	return current;
+}
+
+function isPlainObject(value: unknown): value is Record<string, unknown> {
+	const proto = Object.getPrototypeOf(value as object);
+	return proto === Object.prototype || proto === null;
+}
+
+function unwrapDeep(value: unknown, seen: Set<object>): unknown {
+	const unwrapped = unwrapProxy(value);
+	if (!unwrapped || typeof unwrapped !== "object") {
+		return unwrapped;
+	}
+	if (seen.has(unwrapped)) {
+		return unwrapped;
+	}
+	seen.add(unwrapped);
+
+	if (Array.isArray(unwrapped)) {
+		for (let i = 0; i < unwrapped.length; i++) {
+			const child = unwrapDeep(unwrapped[i], seen);
+			if (child !== unwrapped[i]) {
+				unwrapped[i] = child;
+			}
+		}
+		return unwrapped;
+	}
+
+	if (unwrapped instanceof Map) {
+		const replacements: [unknown, unknown, unknown][] = [];
+		for (const [key, child] of unwrapped.entries()) {
+			const nextKey = unwrapDeep(key, seen);
+			const nextChild = unwrapDeep(child, seen);
+			if (nextKey !== key || nextChild !== child) {
+				replacements.push([key, nextKey, nextChild]);
+			}
+		}
+		for (const [key, nextKey, nextChild] of replacements) {
+			if (nextKey !== key) {
+				unwrapped.delete(key);
+			}
+			unwrapped.set(nextKey, nextChild);
+		}
+		return unwrapped;
+	}
+
+	if (unwrapped instanceof Set) {
+		const replacements: [unknown, unknown][] = [];
+		for (const child of unwrapped.values()) {
+			const next = unwrapDeep(child, seen);
+			if (next !== child) {
+				replacements.push([child, next]);
+			}
+		}
+		for (const [child, next] of replacements) {
+			unwrapped.delete(child);
+			unwrapped.add(next);
+		}
+		return unwrapped;
+	}
+
+	if (isPlainObject(unwrapped)) {
+		for (const key of Object.keys(unwrapped)) {
+			const child = unwrapDeep(unwrapped[key], seen);
+			if (child !== unwrapped[key]) {
+				unwrapped[key] = child;
+			}
+		}
+	}
+
+	return unwrapped;
+}
+
+/**
+ * Strips every `@rivetkit/on-change` proxy out of a value in place, including
+ * proxies nested inside plain objects, arrays, `Map`s, and `Set`s.
+ *
+ * A read of `c.state` or `conn.state` hands back a deep write-through proxy,
+ * so an update written as `c.state = { ...c.state, foo }` produces a plain root
+ * object whose children are still proxies. Persisting that value as-is makes
+ * the next read wrap proxies in another proxy layer, and each layer multiplies
+ * the work of traversing the state, so repeated spread updates degrade
+ * exponentially. Unwrapping before persisting keeps stored state proxy-free.
+ */
+export function unwrapWriteThroughProxy<T>(value: T): T {
+	return unwrapDeep(value, new Set()) as T;
+}

rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-state.test.ts

@@ -54,6 +54,41 @@ describeDriverMatrix("Actor Conn State", (driverTestConfig) => {
 			});
 		});
 
+		describe("State Proxy Nesting", () => {
+			test("does not stack proxies on nested connection state", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				const connection = client.connStateActor
+					.getOrCreate()
+					.connect();
+
+				const result = await connection.spreadUpdateConnState(12);
+
+				// Values read off the state proxy are wrapped exactly once.
+				// Extra layers mean each update persisted the previous read's
+				// proxies, which makes state traversal cost grow exponentially.
+				expect(result.depth).toBe(1);
+				expect(result.tags).toEqual(["read", "write"]);
+
+				await connection.dispose();
+			});
+
+			test("does not stack proxies on nested actor state", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				const connection = client.connStateActor
+					.getOrCreate()
+					.connect();
+
+				const result = await connection.spreadUpdateActorState(12);
+
+				expect(result.depth).toBe(1);
+				expect(result.tags).toEqual(["read", "write"]);
+
+				await connection.dispose();
+			});
+		});
+
 		describe("Connection State Management", () => {
 			test("should maintain unique state for each connection", async (c) => {
 				const { client } = await setupDriverTest(c, driverTestConfig);