fix(rivetkit): normalize waitUntil/keepAwake promises to null to avoid serde undefined error

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts

@@ -167,3 +167,22 @@ export const rawWebSocketBinaryActor = actor({
 	},
 	actions: {},
 });
+
+export const rawWebSocketAsyncOpenActor = actor({
+	state: {
+		openCount: 0,
+	},
+	async onWebSocket(ctx, websocket) {
+		ctx.state.openCount += 1;
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		websocket.send(
+			JSON.stringify({
+				type: "async-open",
+				openCount: ctx.state.openCount,
+			}),
+		);
+	},
+	actions: {
+		getOpenCount: (ctx) => ctx.state.openCount,
+	},
+});

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts

@@ -71,7 +71,11 @@ import {
 	rawHttpVoidReturnActor,
 } from "./raw-http";
 import { rawHttpRequestPropertiesActor } from "./raw-http-request-properties";
-import { rawWebSocketActor, rawWebSocketBinaryActor } from "./raw-websocket";
+import {
+	rawWebSocketActor,
+	rawWebSocketAsyncOpenActor,
+	rawWebSocketBinaryActor,
+} from "./raw-websocket";
 import { rejectConnectionActor } from "./reject-connection";
 import { requestAccessActor } from "./request-access";
 import {
@@ -96,6 +100,7 @@ import {
 	sleepWithRawHttp,
 	sleepWithRawWebSocket,
 	sleepWithWaitUntilMessage,
+	counterWaitUntilProbe,
 	sleepRawWsOnClose,
 	sleepRawWsOnMessage,
 	sleepRawWsSendOnSleep,
@@ -196,6 +201,7 @@ export const registry = setup({
 		sleepRawWsSendOnSleep,
 		sleepRawWsDelayedSendOnSleep,
 		sleepWithWaitUntilInOnWake,
+		counterWaitUntilProbe,
 		// From sleep-db.ts
 		sleepWithDb,
 		sleepWithSlowScheduledDb,
@@ -260,6 +266,7 @@ export const registry = setup({
 		rawHttpRequestPropertiesActor,
 		// From raw-websocket.ts
 		rawWebSocketActor,
+		rawWebSocketAsyncOpenActor,
 		rawWebSocketBinaryActor,
 		// From reject-connection.ts
 		rejectConnectionActor,

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep.ts

@@ -464,6 +464,41 @@ export const sleepWithWaitUntilInOnWake = actor({
 	},
 });
 
+export const counterWaitUntilProbe = actor({
+	state: { count: 0 },
+	actions: {
+		triggerWaitUntilVoid: (c) => {
+			c.state.count += 1;
+			c.waitUntil(Promise.resolve());
+			return c.state.count;
+		},
+		triggerWaitUntilWithValue: (c) => {
+			c.state.count += 1;
+			c.waitUntil(Promise.resolve({ ok: true }));
+			return c.state.count;
+		},
+		triggerWaitUntilRejectVoid: (c) => {
+			c.state.count += 1;
+			c.waitUntil(Promise.reject(new Error("reject-with-error-ok")));
+			return c.state.count;
+		},
+		triggerKeepAwakeVoid: async (c) => {
+			c.state.count += 1;
+			await c.keepAwake(Promise.resolve());
+			return c.state.count;
+		},
+		triggerKeepAwakeWithValue: async (c) => {
+			c.state.count += 1;
+			await c.keepAwake(Promise.resolve({ ok: true }));
+			return c.state.count;
+		},
+		getCount: (c) => c.state.count,
+	},
+	options: {
+		sleepTimeout: SLEEP_TIMEOUT,
+	},
+});
+
 export const sleepWithNoSleepOption = actor({
 	state: { startCount: 0, sleepCount: 0 },
 	onWake: (c) => {
@@ -485,4 +520,3 @@ export const sleepWithNoSleepOption = actor({
 		noSleep: true,
 	},
 });
-

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

@@ -2159,7 +2159,8 @@ class TrackedNativeWebSocketAdapter implements UniversalWebSocket {
 					})
 					.finally(() => {
 						this.#ctx.endWebSocketCallback(callbackRegionId);
-					}),
+					})
+					.then(() => null),
 			);
 		} catch (error) {
 			logger().error({
@@ -2638,8 +2639,9 @@ export class NativeActorContextAdapter {
 	}
 
 	waitUntil(promise: Promise<unknown>): void {
+		const trackedPromise = Promise.resolve(promise).then(() => null);
 		try {
-			callNativeSync(() => this.#ctx.waitUntil(Promise.resolve(promise)));
+			callNativeSync(() => this.#ctx.waitUntil(trackedPromise));
 		} catch (error) {
 			if (!isClosedTaskRegistrationError(error)) {
 				throw error;

rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep.test.ts

@@ -304,6 +304,62 @@ describeDriverMatrix("Actor Sleep", (driverTestConfig) => {
 			}
 		});
 
+		test("waitUntil accepts promises that resolve to undefined", async (c) => {
+			const { client, getRuntimeOutput } = await setupDriverTest(
+				c,
+				driverTestConfig,
+			);
+
+			const probe = client.counterWaitUntilProbe.getOrCreate();
+
+			expect(await probe.triggerWaitUntilVoid()).toBe(1);
+			await waitFor(driverTestConfig, 50);
+
+			expect(await probe.getCount()).toBe(1);
+			expect(getRuntimeOutput()).not.toContain(
+				"undefined cannot be represented as a serde_json::Value",
+			);
+
+			expect(await probe.triggerWaitUntilWithValue()).toBe(2);
+			await waitFor(driverTestConfig, 50);
+
+			expect(await probe.getCount()).toBe(2);
+			expect(getRuntimeOutput()).not.toContain(
+				"undefined cannot be represented as a serde_json::Value",
+			);
+
+			expect(await probe.triggerWaitUntilRejectVoid()).toBe(3);
+			await waitFor(driverTestConfig, 50);
+
+			const runtimeOutput = getRuntimeOutput();
+			expect(runtimeOutput).toContain("actor wait_until promise rejected");
+			expect(runtimeOutput).toContain("reject-with-error-ok");
+			expect(runtimeOutput).not.toContain(
+				"undefined cannot be represented as a serde_json::Value",
+			);
+		});
+
+		test("keepAwake accepts promises that resolve to undefined", async (c) => {
+			const { client, getRuntimeOutput } = await setupDriverTest(
+				c,
+				driverTestConfig,
+			);
+
+			const probe = client.counterWaitUntilProbe.getOrCreate();
+
+			expect(await probe.triggerKeepAwakeVoid()).toBe(1);
+			expect(await probe.triggerKeepAwakeWithValue()).toBe(2);
+			expect(await probe.getCount()).toBe(2);
+
+			const runtimeOutput = getRuntimeOutput();
+			expect(runtimeOutput).not.toContain(
+				"keepAwake bridge to native runtime failed",
+			);
+			expect(runtimeOutput).not.toContain(
+				"undefined cannot be represented as a serde_json::Value",
+			);
+		});
+
 		test("rpc calls keep actor awake", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
 			const actorKey = [`rpc-awake-${crypto.randomUUID()}`];

rivetkit-typescript/packages/rivetkit/tests/driver/raw-websocket.test.ts

@@ -428,6 +428,30 @@ describeDriverMatrix("Raw Websocket", (driverTestConfig) => {
 			expect(finalStats?.connectionCount).toBe(0);
 		});
 
+		test("should handle async onWebSocket open handler", async (c) => {
+			const { client, getRuntimeOutput } = await setupDriverTest(
+				c,
+				driverTestConfig,
+			);
+			const actor = client.rawWebSocketAsyncOpenActor.getOrCreate([
+				"async-open",
+			]);
+
+			const ws = await actor.webSocket();
+			const message = await waitForJsonMessage(ws, 5_000);
+
+			expect(message).toEqual({
+				type: "async-open",
+				openCount: 1,
+			});
+			expect(await actor.getOpenCount()).toBe(1);
+			expect(getRuntimeOutput()).not.toContain(
+				"undefined cannot be represented as a serde_json::Value",
+			);
+
+			ws.close();
+		});
+
 		test("should properly handle onWebSocket open and close events", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
 			const actor = client.rawWebSocketActor.getOrCreate([

rivetkit-typescript/packages/rivetkit/tests/driver/shared-harness.ts

@@ -602,6 +602,7 @@ export async function startNativeDriverRuntime(
 		endpoint,
 		namespace,
 		runnerName: poolName,
+		getRuntimeOutput: () => childOutput(logs),
 		cleanup: async () => {
 			await stopRuntime(runtime);
 		},

rivetkit-typescript/packages/rivetkit/tests/driver/shared-types.ts

@@ -17,6 +17,7 @@ export interface DriverDeployOutput {
 	runnerName: string;
 	hardCrashActor?: (actorId: string) => Promise<void>;
 	hardCrashPreservesData?: boolean;
+	getRuntimeOutput?: () => string;
 	cleanup(): Promise<void>;
 }
 

rivetkit-typescript/packages/rivetkit/tests/driver/shared-utils.ts

@@ -31,6 +31,7 @@ export async function setupDriverTest(
 	endpoint: string;
 	hardCrashActor?: (actorId: string) => Promise<void>;
 	hardCrashPreservesData: boolean;
+	getRuntimeOutput: () => string;
 }> {
 	if (!driverTestConfig.useRealTimers) {
 		vi.useFakeTimers();
@@ -46,6 +47,7 @@ export async function setupDriverTest(
 		runnerName,
 		hardCrashActor,
 		hardCrashPreservesData,
+		getRuntimeOutput,
 		cleanup,
 	} = await driverTestConfig.start();
 	timing("setup.driver_start", driverStartStartedAt, testName);
@@ -83,6 +85,7 @@ export async function setupDriverTest(
 		endpoint,
 		hardCrashActor,
 		hardCrashPreservesData: hardCrashPreservesData ?? false,
+		getRuntimeOutput: getRuntimeOutput ?? (() => ""),
 	};
 }