fix(rivetkit): drain native sleep side tasks reliably

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit-napi/index.d.ts

@@ -225,7 +225,7 @@ export declare class ActorContext {
   disconnectConn(id: string): Promise<void>
   disconnectConns(predicate: (...args: any[]) => any): Promise<void>
   broadcast(name: string, args: Buffer): void
-  waitUntil(promise: Promise<any>): Promise<void>
+  waitUntil(promise: Promise<any>): void
   keepAwake(promise: Promise<any>): Promise<any>
   registerTask(promise: Promise<any>): void
   runtimeState(): object

rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs

@@ -614,7 +614,7 @@ impl ActorContext {
 	}
 
 	#[napi]
-	pub async fn wait_until(&self, promise: Promise<serde_json::Value>) -> napi::Result<()> {
+	pub fn wait_until(&self, promise: Promise<serde_json::Value>) -> napi::Result<()> {
 		self.inner.wait_until(async move {
 			if let Err(error) = promise.await {
 				tracing::warn!(?error, "actor wait_until promise rejected");

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

@@ -146,6 +146,7 @@ type NativeActorRuntimeState = {
 	vars?: unknown;
 	destroyGate?: NativeDestroyGate;
 	persistState?: NativePersistActorState;
+	preventSleepRelease?: () => void;
 };
 
 // Keep JS-only actor caches on the NAPI ActorContext runtime-state bag instead
@@ -2639,20 +2640,13 @@ export class NativeActorContextAdapter {
 	}
 
 	waitUntil(promise: Promise<unknown>): void {
-		// Same counter-arm race as `keepAwake`: increment of the
-		// shutdown_counter happens on first poll of the Rust future. Acceptable
-		// because the only consumer is the grace-finalize predicate, which
-		// debounces through `activity_notify` and re-checks the counter.
-		callNative(() => this.#ctx.waitUntil(Promise.resolve(promise))).catch(
-			(error) => {
-				if (!isClosedTaskRegistrationError(error)) {
-					logger().warn({
-						msg: "waitUntil bridge to native runtime failed",
-						error: stringifyError(error),
-					});
-				}
-			},
-		);
+		try {
+			callNativeSync(() => this.#ctx.waitUntil(Promise.resolve(promise)));
+		} catch (error) {
+			if (!isClosedTaskRegistrationError(error)) {
+				throw error;
+			}
+		}
 	}
 
 	beginWebSocketCallback(): number {
@@ -2665,12 +2659,30 @@ export class NativeActorContextAdapter {
 		);
 	}
 
-	/** @deprecated No-op. Use `keepAwake(promise)` or `waitUntil(promise)` instead. */
-	setPreventSleep(_preventSleep: boolean): void {}
+	/** @deprecated Use `keepAwake(promise)` or `waitUntil(promise)` instead. */
+	setPreventSleep(preventSleep: boolean): void {
+		const runtimeState = getNativeRuntimeState(this.#ctx);
+		if (preventSleep) {
+			if (runtimeState.preventSleepRelease) {
+				return;
+			}
+
+			let release: () => void = () => {};
+			const promise = new Promise<void>((resolve) => {
+				release = resolve;
+			});
+			runtimeState.preventSleepRelease = release;
+			this.keepAwake(promise);
+			return;
+		}
+
+		runtimeState.preventSleepRelease?.();
+		runtimeState.preventSleepRelease = undefined;
+	}
 
-	/** @deprecated No-op. Always returns `false`. */
+	/** @deprecated Use `keepAwake(promise)` or `waitUntil(promise)` instead. */
 	get preventSleep(): boolean {
-		return false;
+		return getNativeRuntimeState(this.#ctx).preventSleepRelease !== undefined;
 	}
 
 	sleep(): void {
@@ -3743,10 +3755,9 @@ export function buildNativeFactory(
 							try {
 								if (typeof config.onSleep === "function") {
 									await config.onSleep(actorCtx);
+									await actorCtx.saveState({ immediate: true });
 								}
 							} finally {
-								await actorCtx.closeDatabase(false);
-								callNativeSync(() => ctx.clearRuntimeState());
 								await actorCtx.dispose();
 							}
 						},
@@ -3766,7 +3777,6 @@ export function buildNativeFactory(
 							} finally {
 								resolveNativeDestroy(ctx);
 								await actorCtx.closeDatabase(true);
-								callNativeSync(() => ctx.clearRuntimeState());
 								await actorCtx.dispose();
 							}
 						},