fix(rivetkit): drain native sleep side tasks reliably

[NathanFlurry]

Description

Please include a summary of the changes and the related issue. Please also include relevant motivation and context.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

How Has This Been Tested?

Please describe the tests that you ran to verify your changes.

Checklist:

• My code follows the style guidelines of this project
• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[NathanFlurry]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• test(rivetkit): avoid polling sleep actor awake #4783
• fix(rivetkit): drain native sleep side tasks reliably #4782 👈 (View in Graphite)
• fix(rivetkit): route raw request fetches to actors #4781
• fix(rivetkit): run db migrations before lifecycle hooks #4780
• fix(sqlite): tolerate startup takeover races #4779
• fix(pegboard): refresh runner config after envoy connect #4778
• fix(gasoline): skip corrupt active worker indexes #4777
• chore(kitchen-sink): configurable endpoint #4766 : 1 other dependent PR (#4754 )
• refactor(rivetkit-core): rename StopReason to ShutdownKind and kill enum fallthroughs #4765
• chore(envoy-client): log command and event payloads #4751
• fix(rivetkit): minor sleep-cleanup follow-ups #4761
• docs(rivetkit): clean up stale setPreventSleep references in docs and JSDoc #4760
• docs(rivetkit): document keepAwake/waitUntil counter-arm race and surface bridge errors #4759
• test(engine): force overdue branch in alarm-during-sleep test with negative offset #4758
• fix(rivetkit-core): return stopping not starting when sleep/destroy called mid-shutdown #4757
• test(rivetkit): remove sleepWithPreventSleep fixture and tests #4756
• test(engine): add regression test for alarm-during-sleep-transition race #4755
• test(rivetkit-core): cover sleep/destroy guards and prevent_sleep no-op #4750
• docs: add keepAwake vs waitUntil table and internal sleep-sequence invariants #4749
• feat(rivetkit): expose c.keepAwake(promise) through native path #4748
• refactor(rivetkit): deprecate setPreventSleep to no-op in TypeScript #4747
• fix(rivetkit-core): enumerate non-drained counters in grace deadline warn log #4746
• fix(rivetkit-core): abort graceful cleanup task when grace deadline elapses #4745
• fix(rivetkit-core): error on sleep/destroy before startup or already-requested #4743
• fix(rivetkit-core): cap SerializeState shutdown timeout at 15s #4742
• refactor(rivetkit-core): remove prevent_sleep field and deprecate set_prevent_sleep to no-op #4741
• fix(api-public): remove orphan import_export actors module #4764
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[claude]

PR Review: fix(rivetkit): drain native sleep side tasks reliably

Summary

This PR fixes waitUntil to use a synchronous bridge call instead of an async fire-and-forget, and re-implements setPreventSleep/preventSleep as a real keepAwake-backed gate. The overall direction is good — the synchronous bridge is cleaner and avoids the async race in the old callNative(...).catch(...) pattern.

---

Bug: closeDatabase(false) missing from onSleep finally block

The onSleep handler now only calls actorCtx.closeDatabase(false) inside the if (typeof config.onSleep === "function") branch, but the outer condition already gates on typeof config.onSleep === "function" || databaseProvider !== undefined.

When databaseProvider is set but config.onSleep is not defined, the code goes straight to finally { await actorCtx.dispose() } without ever calling closeDatabase. This is a regression — the database provider's onDestroy cleanup never runs for sleep/wake cycles.

Additionally, the saveState({ immediate: true }) call is also inside the config.onSleep branch, so the database-provider-only path skips state persistence too.

The finally block should unconditionally close the database:

} finally {
    await actorCtx.closeDatabase(false);
    await actorCtx.dispose();
}

---

Observation: clearRuntimeState removal

clearRuntimeState() was previously called in both onSleep and onDestroy finally blocks to explicitly unref the JS runtime state bag from the Rust shared context. Its removal means cleanup now depends on GC rather than being deterministic.

The preventSleepRelease closure added by this PR lives on that same runtime state bag. Once the promise is resolved, the closure is benign, but losing the deterministic clearRuntimeState call removes the explicit GC-root release that the old code relied on. Worth confirming this is intentional.

---

Observation: keepAwake async race still present in setPreventSleep

setPreventSleep(true) calls this.keepAwake(promise), which still uses the async callNative(...) bridge internally. The acknowledged sub-millisecond window where keep_awake_count == 0 could be observed still exists in this path. This is pre-existing and documented, but worth noting since the PR's primary goal is draining side tasks reliably.

---

Style: @deprecated annotation on setPreventSleep

The JSDoc still marks setPreventSleep as @deprecated, but the function now has a real implementation. If this is intended as a legacy compatibility shim, the annotation should note that it is a shim backed by keepAwake. If it is no longer deprecated, the annotation should be removed.

---

Code quality

The waitUntil change from async to synchronous is clean and correct. The callNativeSync + try/catch pattern matches the established convention for bridge calls that should not be fire-and-forget. The index.d.ts update to void return type is consistent.

The setPreventSleep re-implementation with a pending promise and keepAwake is a reasonable approach for backward compatibility. The stored preventSleepRelease on NativeActorRuntimeState is a clear pattern.

---

Verdict

The closeDatabase(false) omission in the database-provider-only sleep path is the main issue to address before merging. The rest are observations worth considering but not blockers.

[github-actions]

Preview packages published to npm

Install with:

npm install rivetkit@pr-4782

All packages published as 0.0.0-pr.4782.868674a with tag pr-4782.

Engine binary is shipped via @rivetkit/engine-cli on linux-x64-musl, linux-arm64-musl, darwin-x64, and darwin-arm64. Windows users should use the release installer or set RIVET_ENGINE_BINARY.

Docker images:

docker pull rivetdev/engine:slim-868674a
docker pull rivetdev/engine:full-868674a

Individual packages

npm install rivetkit@pr-4782
npm install @rivetkit/react@pr-4782
npm install @rivetkit/rivetkit-napi@pr-4782
npm install @rivetkit/workflow-engine@pr-4782