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]

• fix(rivetkit): drain native sleep side tasks reliably #4782 : 2 dependent PRs (#4783 , #4784 ) 👈 (View in Graphite)
• fix(rivetkit): route raw request fetches to actors #4781
• fix(pegboard): refresh runner config after envoy connect #4778 : 2 other dependent PRs (#4779 , #4798 )
• chore(kitchen-sink): configurable endpoint #4766 : 2 other dependent PRs (#4754 , #4777 )
• 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

The core correctness fix — making waitUntil registration synchronous — is solid and closes a real race where the WaitUntil shutdown counter could be evaluated before tasks were registered, letting the actor sleep with in-flight work undrained. Several bundled changes (database lifecycle simplification, setPreventSleep deprecation) deserve individual scrutiny.

---

Core Fix: waitUntil Made Synchronous ✅

The change from pub async fn wait_until to pub fn wait_until (and callNative → callNativeSync in JS) is the right fix. track_shutdown_task is now called immediately, incrementing shutdown_task_count before waitUntil returns, which correctly gates can_finalize_sleep. The old code acknowledged this race: "increment of the shutdown_counter happens on first poll of the Rust future."

The .d.ts change (Promise<void> to void) correctly matches the new Rust signature. Rethrowing non-task-registration errors instead of silently logging them aligns with the fail-by-default rule.

---

Issues

1. setPreventSleep/preventSleep warnings will spam logs for agent-os

The two previously-silent no-ops now emit logger().warn(...) on every call. The agent-os actor calls syncPreventSleep across at least 10 call sites in agent-os/actor/{process,shell,session,index}.ts — on every process start/stop, session event, hook run, and shell operation. This will generate per-call warning spam in production.

Options:

• Emit the warning only once per actor instance via a flag in runtime state
• Rely on the @deprecated JSDoc annotation alone and remove the runtime warn calls
• Update agent-os in this same PR to use waitUntil/keepAwake so no warnings fire at all

2. Narrowed onSleep condition may leave JS database wrapper un-cleared between sleep/wake

Old: typeof config.onSleep === 'function' || databaseProvider !== undefined
New: typeof config.onSleep === 'function'

Actors with a databaseProvider but no explicit onSleep no longer get the JS-side closeDatabase(false) cleanup on sleep. The underlying Rust VFS handle is cleaned up by core, and onMigrate resets the client on wake — so this is functionally safe. However, runtimeState.databaseClient is not cleared between sleep and the next wake's onMigrate, leaving a stale wrapper pointing at a closed native handle for that window.

3. saveState placement in onSleep callback

If config.onSleep throws, saveState is skipped (it is after the onSleep call inside try, not in finally). Core runs save_final_state() unconditionally so state is not lost, but this looks like an unintentional double-write in the normal path and a silent skip in the error path. Either move saveState to finally, or remove it entirely and rely on core's save_final_state().

4. JS runtime state unref removed — leak grows with sleep cycles

The old code called callNativeSync(() => ctx.clearRuntimeState()) in both sleep and destroy paths. The new code removes these, relying on reset_runtime_shared_state() at wake time via mem::forget. For long-running hosts with many sleep cycles the leak grows linearly. If this tradeoff is intentional, document it at the mem::forget site.

---

Positive Changes

• Core waitUntil race fix is correct and well-targeted
• Error rethrowing in wait_until matches fail-by-default
• DatabaseProvider.onDestroy removal simplifies lifecycle (@experimental makes the breaking change acceptable)
• provider field removal from NativeDatabaseClientState is a clean simplification
• Test cleanup is consistent with removing DatabaseProvider.onDestroy

---

Minor

PR description is empty and checklist items are unchecked. The behavioral changes (setPreventSleep warnings, narrowed onSleep condition, new saveState call) should be described so reviewers can assess intent vs. implementation.

[github-actions]

Preview packages published to npm

Install with:

npm install rivetkit@pr-4782

All packages published as 0.0.0-pr.4782.2c31448 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-2c31448
docker pull rivetdev/engine:full-2c31448

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