docs: add keepAwake vs waitUntil table and internal sleep-sequence invariants

Author: NathanFlurry

CLAUDE.md

@@ -425,6 +425,7 @@ Load these only when the task touches the topic.
 - **[BARE protocol crates](docs-internal/engine/bare-protocol-crates.md)** — vbare schema ordering, identity converters, `build.rs` TS codec generation pattern. Read before adding/changing protocol crates.
 - **[SQLite VFS parity](docs-internal/engine/sqlite-vfs.md)** — native Rust VFS ↔ WASM TypeScript VFS 1:1 parity rule, v2 storage keys, chunk layout, delete/truncate strategy. Read before touching either VFS.
 - **[TLS trust roots](docs-internal/engine/tls-trust-roots.md)** — rustls native+webpki union rationale, which clients use which backend.
+- **[Sleep sequence](docs-internal/engine/sleep-sequence.md)** — engine lifecycle authority, `keepAwake` vs `waitUntil` semantics, grace deadline shutdown-token abort, `can_arm_sleep_timer` vs `can_finalize_sleep` predicates. Read before touching sleep/destroy lifecycle.
 
 ### Agent procedural (`.claude/reference/`)
 

docs-internal/engine/sleep-sequence.md

@@ -0,0 +1,64 @@
+# Sleep sequence invariants
+
+Design constraints and invariants for the RivetKit actor sleep / destroy lifecycle. Pair with `actor-task-dispatch.md` and `rivetkit-core-internals.md` for surrounding context.
+
+## Authority
+
+- The engine owns lifecycle authority. `ctx.sleep()` and `ctx.destroy()` send fire-and-forget `ActorIntent` events; they do not transition lifecycle state locally. The local `SleepGrace` / `DestroyGrace` transition runs when the engine replies with `StopActor`.
+- `envoy-client` retries intent delivery across reconnects via checkpoint-based event replay (`engine/sdks/rust/envoy-client/src/events.rs`). Core does not need its own retry path.
+
+## Public surface: keep-awake primitives
+
+Two user-facing primitives in TypeScript. Both accept a `Promise`, never a closure.
+
+| Method | Blocks idle sleep | Blocks grace finalize | Notes |
+| --- | --- | --- | --- |
+| `c.keepAwake(promise)` | Yes | Yes | Returns the same promise. Use for work the actor must stay up for. |
+| `c.waitUntil(promise)` | No | Yes | Returns void. Use for best-effort flush/cleanup work that is allowed to complete inside the grace window. |
+
+`c.setPreventSleep(b)` and `c.preventSleep` are deprecated no-ops retained for binary / call-site compatibility. They will be removed in 2.2.0.
+
+### Why two primitives and not one
+
+`keepAwake` is scoped, non-leaky, and symmetric with `waitUntil`. `setPreventSleep` was a flag that had to be paired by hand; forgetting to clear it wedged the actor awake. A promise-scoped counter cannot leak: when the promise settles (resolve or reject), the counter decrements.
+
+### Why separate `keep_awake` and `internal_keep_awake` in core
+
+Kept separate for debug visibility. Grace deadline warn logs report each counter independently so diagnostics distinguish user keep-awake sites from framework-owned keep-awake sites (schedule alarms, queue receives).
+
+## Sleep readiness predicates
+
+Two predicates govern the sleep state machine. Both live on `ActorContext` / `SleepState`.
+
+- `can_arm_sleep_timer()` — the idle predicate. Returns `CanSleep::Yes` only when every sleep-affecting counter is zero and the run handler is inactive (or waiting on a queue). Used to start the sleep idle timer.
+- `can_finalize_sleep()` — the grace predicate. Returns `true` only when every shutdown-affecting counter is zero: `core_dispatched_hooks`, `shutdown_task_count`, `sleep_keep_awake`, `sleep_internal_keep_awake`, `active_http_requests`, `websocket_callbacks`, `pending_disconnects`. Used to advance from `SleepGrace` to `SleepFinalize` (or finalize destroy).
+
+Removing `preventSleep` deleted both predicate branches. Any future sleep-affecting counter must add an entry in each predicate and must call `ActorContext::reset_sleep_timer()` on transitions that change the result.
+
+## Grace period and abort signals
+
+- `start_grace(reason)` fires at the start of `SleepGrace` / `DestroyGrace`. It cancels the sleep idle timer, cancels the actor abort signal (`actor_abort_signal`), installs a `SleepGraceState` with the effective grace deadline, and resets the sleep timer to arm the grace tick.
+- The actor abort signal is a soft signal: "shutdown has started, please wrap up." User code observes it via `c.abortSignal`. It does not force-stop work.
+- For destroy, the abort signal may fire earlier than grace entry because `ctx.destroy()` cancels the abort token immediately via `mark_destroy_requested(...)`.
+
+## Grace deadline enforcement
+
+When the grace deadline elapses before `can_finalize_sleep()` returns true:
+
+- `on_sleep_grace_deadline` aborts the user `run` handle (`run_handle.abort()`), cancels the shutdown deadline token (`cancel_shutdown_deadline()`), records the timeout, and emits a structured warn log enumerating every non-drained counter.
+- The NAPI `RunGracefulCleanup` task observes `shutdown_deadline_token()` via `tokio::select!` and aborts its in-flight `onSleep` / `onDestroy` call so SQLite and KV cleanup in `teardown_sleep_state` do not race against mid-commit user work.
+- Foreign-runtime adapters that run user cleanup callbacks must observe the shutdown deadline token the same way.
+
+## Guarding lifecycle requests
+
+- `ctx.sleep()` and `ctx.destroy()` return `Result<()>`. They fail with `actor/starting` if called before startup completes and `actor/stopping` if the request flag has already been swapped to true for this generation. An atomic `swap(true, ...)` on `sleep_requested` / `destroy_requested` enforces single-shot request semantics per generation.
+- The idle sleep timer request path (`spawn_sleep_timer_task`) and the `ActorTask` sleep-tick path both suppress the already-requested error: idle-driven requests may race user-driven requests and the warning is informational.
+
+## Serialize-state shutdown cap
+
+`SERIALIZE_STATE_SHUTDOWN_SANITY_CAP = 15s` is the upper bound on how long the shutdown `SerializeState` reply wait is allowed to pend before `save_final_state` falls back to empty deltas (preserving prior state). This is a sanity cap, not a deadline anyone should ever hit; the normal drain finishes in milliseconds.
+
+## Test harness parity
+
+- Rust integration tests live in `rivetkit-core/tests/modules/sleep.rs` and pin predicate behavior, grace period selection, and `save_final_state` cap.
+- TypeScript driver tests in `rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep*.test.ts` cover abort-signal-at-grace-entry, `keepAwake` holding shutdown, `c.db` writes surviving `onSleep`, and regression coverage for `setPreventSleep` being a no-op.

rivetkit-rust/packages/rivetkit-core/CLAUDE.md

@@ -6,7 +6,10 @@
 
 ## Sleep state invariants
 
-- Any mutation that changes a `can_sleep` input must call `ActorContext::reset_sleep_timer()` so the `ActorTask` sleep deadline is re-evaluated. Inputs are: `ready`/`started`, `prevent_sleep`, `no_sleep`, `active_http_request_count`, `sleep_keep_awake_count`, `sleep_internal_keep_awake_count`, `pending_disconnect_count`, `conns()`, and `websocket_callback_count`. Missing this call leaves the sleep timer armed against stale state and triggers the `"sleep idle deadline elapsed but actor stayed awake"` warning on the next tick.
+- Any mutation that changes a `can_sleep` input must call `ActorContext::reset_sleep_timer()` so the `ActorTask` sleep deadline is re-evaluated. Inputs are: `ready`/`started`, `no_sleep`, `active_http_request_count`, `sleep_keep_awake_count`, `sleep_internal_keep_awake_count`, `pending_disconnect_count`, `conns()`, and `websocket_callback_count`. Missing this call leaves the sleep timer armed against stale state and triggers the `"sleep idle deadline elapsed but actor stayed awake"` warning on the next tick.
+- `ActorContext::set_prevent_sleep(...)` / `prevent_sleep()` are deprecated no-ops kept for NAPI bridge compatibility. Use `keep_awake(future)` (holds counter while awaited) or `wait_until(future)` (tracked shutdown task) instead. Do not reintroduce a `prevent_sleep` field, a `CanSleep::PreventSleep` variant, or branches that read it.
+- `ctx.sleep()` and `ctx.destroy()` return `Result<()>`. They error with `ActorLifecycleError::Starting` when called before startup completes and `ActorLifecycleError::Stopping` if the requested flag has already been set this generation (atomic `swap(true, ...)`). Internal idle-timer paths log and suppress the already-requested error.
+- The grace deadline path (`on_sleep_grace_deadline`) aborts the user `run` handle and cancels `shutdown_deadline_token()`. Foreign-runtime adapters running `onSleep` / `onDestroy` must observe that token via `tokio::select!` so SQLite teardown does not race user cleanup work.
 - Counter `register_zero_notify(&idle_notify)` hooks only drive shutdown drain waits. They are not a substitute for the activity-dirty notification, so any new sleep-affecting counter must also notify on transitions that change `can_sleep`.
 - A clean `run` exit while `Started` is not terminal. Keep the generation alive until the guaranteed `Stop` drives `SleepGrace` or `DestroyGrace`, and only treat `Terminated` as "grace hooks already completed."
 - Do not reply to actor startup until the runtime adapter has acknowledged its startup preamble. Otherwise `getOrCreate` can race the first action against `onWake` or `run` startup.

website/src/content/docs/actors/lifecycle.mdx

@@ -285,7 +285,7 @@ The `run` hook is called after the actor starts and runs in the background witho
 The handler exposes `c.aborted` for loop checks and `c.abortSignal` for canceling operations when the actor is stopping. You should always check or listen for shutdown to exit gracefully.
 
 **Important behavior:**
-- The actor may go to sleep at any time during the `run` handler. Use `c.setPreventSleep(true)` while work is active, then clear it with `c.setPreventSleep(false)` once the actor can sleep again.
+- The actor may go to sleep at any time during the `run` handler. Wrap work that must keep the actor awake with `c.keepAwake(promise)` to block idle sleep until the promise settles.
 - If the `run` handler exits (returns), the actor follows its normal idle sleep timeout once it becomes idle
 - If the `run` handler throws an error, the actor logs the error and then follows its normal idle sleep timeout once it becomes idle
 - On shutdown, `c.abortSignal` fires so the `run` handler can exit within the graceful shutdown window.
@@ -725,7 +725,7 @@ An actor is considered idle and eligible to sleep when **all** of the following
 - No active HTTP requests
 - No active connections (unless they are hibernatable WebSockets)
 - No active `run` handler (unless it is waiting on a queue)
-- `setPreventSleep` is not enabled
+- No outstanding `c.keepAwake(promise)` promises
 - No pending disconnect callbacks
 - No async `onWebSocket` event handlers (eg `open`, `message`, `close`) still running
 
@@ -737,13 +737,18 @@ Outbound requests (e.g. `fetch` calls) do not count as activity and will not kee
 
 The platform may force an actor to migrate to a new machine during version upgrades or when a serverless request is about to timeout. The same [shutdown sequence](#shutdown-sequence) runs, then the actor is rescheduled on a new machine and wakes up with its persisted state.
 
-Use `onSleep`, `waitUntil`, or `setPreventSleep` to control the length of the grace period before the actor moves to another machine.
+Use `onSleep`, `waitUntil`, or `keepAwake` to control the length of the grace period before the actor moves to another machine.
 
-### Preventing Sleep
+### Keeping the Actor Awake
 
-If actor state says the actor should stay awake, call `c.setPreventSleep(true)` and clear it once the actor can sleep again. You can read `c.preventSleep` to inspect the current flag.
+RivetKit gives you two primitives for holding the actor awake across background work. Both take a `Promise` and differ in how they interact with idle sleep and the grace period.
 
-`setPreventSleep` blocks normal idle sleep until you clear it. It is not a platform-wide stop blocker though. If shutdown has already started, RivetKit waits for `preventSleep` to clear within the same `sleepGracePeriod` shutdown budget used by `onSleep` and `waitUntil`.
+| Method | Accepts | Blocks idle sleep | Blocks grace finalize | Use case |
+| --- | --- | --- | --- | --- |
+| `c.keepAwake(promise)` | `Promise<T>` (returns same promise) | Yes | Yes | Critical work that must keep the actor running end to end (for example a turn in a game, an ongoing tool call). |
+| `c.waitUntil(promise)` | `Promise<unknown>` (returns void) | No | Yes | Best-effort finalization work that may complete during the grace window (for example analytics flushes, cleanup writes). |
+
+`c.keepAwake(promise)` is the preferred primitive for long-running work the actor should not sleep through. It holds a keep-awake counter until the promise settles, which blocks both idle sleep and the grace finalize step. The promise is returned unchanged, so you can `await` it if you need the value.
 
 ```typescript
 import { actor } from "rivetkit";
@@ -754,18 +759,25 @@ const sessionActor = actor({
   },
 
   actions: {
-    beginTurn: (c) => {
+    runTurn: async (c, input: string) => {
       c.state.activeTurns += 1;
-      c.setPreventSleep(true);
-    },
-    endTurn: (c) => {
-      c.state.activeTurns -= 1;
-      c.setPreventSleep(c.state.activeTurns > 0);
+      try {
+        const result = await c.keepAwake(processTurn(input));
+        return result;
+      } finally {
+        c.state.activeTurns -= 1;
+      }
     },
   }
 });
+
+declare function processTurn(input: string): Promise<string>;
 ```
 
+<Note>
+`setPreventSleep(enabled)` is deprecated and now a no-op. Wrap the work you want to keep alive with `c.keepAwake(promise)` instead.
+</Note>
+
 ### On Sleep Hook
 
 The [`onSleep`](#onsleep) hook runs during shutdown for cleanup like clearing intervals or closing connections. It is best-effort and will not run if the actor crashes.
@@ -815,14 +827,14 @@ const analyticsActor = actor({
 });
 ```
 
-The actor waits up to `sleepGracePeriod` for graceful sleep work during the [shutdown sequence](#shutdown-sequence). That single budget covers `onSleep`, `waitUntil`, async raw WebSocket handlers such as `message` and `close`, and waiting for `preventSleep` to clear after shutdown has started. By default, this graceful sleep window is 15 seconds total. If the timeout is exceeded, the actor proceeds with sleep anyway.
+The actor waits up to `sleepGracePeriod` for graceful sleep work during the [shutdown sequence](#shutdown-sequence). That single budget covers `onSleep`, `waitUntil`, `keepAwake`, async raw WebSocket handlers such as `message` and `close`. By default, this graceful sleep window is 15 seconds total. If the timeout is exceeded, the actor proceeds with sleep anyway.
 
 ### Sleep Timeouts
 
 | Option | Default | Description |
 |--------|---------|-------------|
 | `sleepTimeout` | 30 seconds | Time of inactivity before the actor begins sleeping. |
-| `sleepGracePeriod` | 15 seconds | Total graceful shutdown window for hooks, `waitUntil`, async raw WebSocket handlers, disconnects, and waiting for `preventSleep` to clear. |
+| `sleepGracePeriod` | 15 seconds | Total graceful shutdown window for hooks, `waitUntil`, `keepAwake`, async raw WebSocket handlers, and disconnects. |
 
 Rivet enforces a hard limit of **30 minutes** for the entire stop process. These can be configured in the [options](#options).