diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-state.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-state.ts index 9725eddb77..77086177bf 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-state.ts +++ b/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); + 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; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index c09df3f690..b2330a5887 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/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; 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 { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/write-through-proxy.ts b/rivetkit-typescript/packages/rivetkit/src/registry/write-through-proxy.ts index 3ceab45f8f..f636958837 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/write-through-proxy.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/write-through-proxy.ts @@ -38,3 +38,104 @@ export function createWriteThroughProxy( }, ) 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); + if (target === current) { + break; + } + current = target; + } + return current; +} + +function isPlainObject(value: unknown): value is Record { + const proto = Object.getPrototypeOf(value as object); + return proto === Object.prototype || proto === null; +} + +function unwrapDeep(value: unknown, seen: Set): 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(value: T): T { + return unwrapDeep(value, new Set()) as T; +} diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-state.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-state.test.ts index 204ee9a251..6fd4573e30 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-state.test.ts +++ b/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); diff --git a/rivetkit-typescript/packages/rivetkit/tests/write-through-proxy.test.ts b/rivetkit-typescript/packages/rivetkit/tests/write-through-proxy.test.ts index ad1c646389..7bf0a22109 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/write-through-proxy.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/write-through-proxy.test.ts @@ -1,10 +1,14 @@ +import onChange from "@rivetkit/on-change"; import { describe, expect, test, vi } from "vitest"; import { assertJsonCompatValue, encodeJsonCompatValue, reviveJsonCompatValue, } from "@/common/encoding"; -import { createWriteThroughProxy } from "@/registry/write-through-proxy"; +import { + createWriteThroughProxy, + unwrapWriteThroughProxy, +} from "@/registry/write-through-proxy"; import { decodeCborCompat, encodeCborCompat } from "@/serde"; describe("createWriteThroughProxy", () => { @@ -282,3 +286,139 @@ describe("encodeJsonCompatValue validation", () => { expect(() => encodeJsonCompatValue(obj)).toThrow(TypeError); }); }); + +/** + * Counts how many `@rivetkit/on-change` proxy layers wrap a value. A raw value + * has depth 0. + */ +function proxyDepth(value: unknown): number { + let depth = 0; + let current = value; + while (current !== null && typeof current === "object") { + const target = onChange.target(current as Record); + if (target === current) { + break; + } + depth++; + current = target; + } + return depth; +} + +describe("unwrapWriteThroughProxy", () => { + test("strips proxies nested in objects, arrays, maps, and sets", () => { + const raw = { + nested: { inner: { deep: true } }, + arr: [{ item: 1 }], + map: new Map([["a", { value: 1 }]]), + set: new Set([{ value: 2 }]), + }; + const proxied = createWriteThroughProxy(raw, () => {}); + + // Reading through the proxy hands back proxies for every child. + const spread = { + nested: proxied.nested, + arr: proxied.arr, + map: proxied.map, + set: proxied.set, + }; + expect(proxyDepth(spread.nested)).toBe(1); + + const unwrapped = unwrapWriteThroughProxy(spread); + expect(proxyDepth(unwrapped.nested)).toBe(0); + expect(proxyDepth(unwrapped.nested.inner)).toBe(0); + expect(proxyDepth(unwrapped.arr)).toBe(0); + expect(proxyDepth(unwrapped.arr[0])).toBe(0); + expect(proxyDepth(unwrapped.map)).toBe(0); + expect(proxyDepth(unwrapped.map.get("a"))).toBe(0); + expect(proxyDepth([...unwrapped.set][0])).toBe(0); + + // Unwrapping preserves structure and identity of the raw tree. + expect(unwrapped.nested).toBe(raw.nested); + expect(unwrapped.map.get("a")).toEqual({ value: 1 }); + expect([...unwrapped.set]).toEqual([{ value: 2 }]); + }); + + test("returns primitives unchanged and tolerates cycles", () => { + expect(unwrapWriteThroughProxy(null)).toBe(null); + expect(unwrapWriteThroughProxy(undefined)).toBe(undefined); + expect(unwrapWriteThroughProxy(7)).toBe(7); + expect(unwrapWriteThroughProxy("x")).toBe("x"); + + const cyclic: Record = { name: "root" }; + cyclic.self = cyclic; + expect(unwrapWriteThroughProxy(cyclic)).toBe(cyclic); + }); + + test("unwraps a value that is itself a proxy", () => { + const raw = { a: 1 }; + const proxied = createWriteThroughProxy(raw, () => {}); + expect(unwrapWriteThroughProxy(proxied)).toBe(raw); + }); +}); + +describe("spread state updates", () => { + // Mirrors how `NativeConnAdapter` and the actor context expose state: the + // getter hands back a deep write-through proxy and the setter persists the + // assigned value. + function createStateHolder(initial: unknown, unwrapOnWrite: boolean) { + let stored = initial; + return { + get raw() { + return stored; + }, + get state(): any { + return createWriteThroughProxy(stored, (next) => { + stored = next; + }); + }, + set state(value: unknown) { + stored = unwrapOnWrite ? unwrapWriteThroughProxy(value) : value; + // Stand-in for the encode and validate traversal both adapters + // run on every write. + assertJsonCompatValue(stored); + encodeCborCompat(stored); + }, + }; + } + + const initialState = () => ({ + capability: "write", + executorCapabilities: { + tags: ["executor-relay-v1", "filesystem-v2-paths-v1"], + }, + }); + + test("stacks proxy layers on nested values when the raw value is stored", () => { + const holder = createStateHolder(initialState(), false); + + for (let i = 0; i < 5; i++) { + holder.state = { ...holder.state, capability: "write" }; + } + + // Each update stores the previous read's proxies, so the next read + // wraps them again. This nesting is what makes state traversal cost + // grow exponentially. + expect( + proxyDepth((holder.raw as any).executorCapabilities), + ).toBeGreaterThan(1); + }); + + test("keeps stored state proxy-free across repeated spread updates", () => { + const holder = createStateHolder(initialState(), true); + + const started = performance.now(); + for (let i = 0; i < 12; i++) { + holder.state = { ...holder.state, capability: "write" }; + } + const elapsed = performance.now() - started; + + expect(proxyDepth((holder.raw as any).executorCapabilities)).toBe(0); + expect((holder.raw as any).executorCapabilities.tags).toEqual([ + "executor-relay-v1", + "filesystem-v2-paths-v1", + ]); + // Without unwrapping, twelve updates of this tiny state take minutes. + expect(elapsed).toBeLessThan(1000); + }); +});