From 7f23a4c3b8af6ebe8a37a35b45598d20ec14da29 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Tue, 28 Apr 2026 03:11:36 -0700 Subject: [PATCH] fix(rivetkit): normalize waitUntil/keepAwake promises to null to avoid serde undefined error --- .../driver-test-suite/raw-websocket.ts | 19 +++++++ .../driver-test-suite/registry-static.ts | 9 ++- .../fixtures/driver-test-suite/sleep.ts | 36 +++++++++++- .../packages/rivetkit/src/registry/native.ts | 6 +- .../rivetkit/tests/driver/actor-sleep.test.ts | 56 +++++++++++++++++++ .../tests/driver/raw-websocket.test.ts | 24 ++++++++ .../rivetkit/tests/driver/shared-harness.ts | 1 + .../rivetkit/tests/driver/shared-types.ts | 1 + .../rivetkit/tests/driver/shared-utils.ts | 3 + 9 files changed, 151 insertions(+), 4 deletions(-) diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts index 696e33d6aa..884ddd83f2 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts +++ b/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, + }, +}); diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts index a06fdcacaa..c58fa23abf 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts +++ b/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, diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep.ts index cf2a499ad3..28da755263 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep.ts +++ b/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, }, }); - diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 7eb3851710..ff8bef74ab 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/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): 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; diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep.test.ts index 39c4492b97..946d363227 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep.test.ts +++ b/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()}`]; diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/raw-websocket.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/raw-websocket.test.ts index 462b72d2b4..3b575b099a 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/raw-websocket.test.ts +++ b/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([ diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/shared-harness.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/shared-harness.ts index bfefb4e5d5..57f6eb8c9b 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/shared-harness.ts +++ b/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); }, diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/shared-types.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/shared-types.ts index e9879a4510..12f6361cfc 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/shared-types.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/driver/shared-types.ts @@ -17,6 +17,7 @@ export interface DriverDeployOutput { runnerName: string; hardCrashActor?: (actorId: string) => Promise; hardCrashPreservesData?: boolean; + getRuntimeOutput?: () => string; cleanup(): Promise; } diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/shared-utils.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/shared-utils.ts index d4cc8f07f3..3ef7faf17b 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/shared-utils.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/driver/shared-utils.ts @@ -31,6 +31,7 @@ export async function setupDriverTest( endpoint: string; hardCrashActor?: (actorId: string) => Promise; 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 ?? (() => ""), }; }