fix(rivetkit): run onSleep after in-flight work and refuse dispatch once it starts

docs-internal/engine/sleep-sequence.md

@@ -40,7 +40,9 @@ Removing `preventSleep` deleted both predicate branches. Any future sleep-affect
 - `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.
 - Destroy requests also use the normal grace path. The actor abort signal fires when destroy grace starts.
-- Sleep grace admits new work only while its hooks still run: the app's `onSleep` and connection disconnects, counted in `core_dispatched_hooks`. After they finish, the app has released its resources (the TypeScript sleep hook also closes SQLite and clears runtime state), so dispatch refuses with `actor/restarting` and the client retries against the next generation. `SleepFinalize` refuses the same way.
+- Sleep grace runs `onSleep` and connection disconnects only after tracked work that drains shutdown grace (actions, `keepAwake`, `waitUntil`, registered tasks, WebSocket callbacks) finishes, or at the hooks deadline, which keeps the last quarter of the grace period for the hooks and the final save. Raw `onRequest` handlers and `run` are not waited for. Destroy grace runs its hooks at once.
+- Sleep grace admits new work until the hooks run. After that, dispatch refuses with `actor/restarting` and the client retries against the next generation. `SleepFinalize` refuses the same way.
+- Core owns SQLite teardown: it syncs the alarm and then closes SQLite in final cleanup. The TypeScript sleep cleanup closes SQLite itself only when tracked work outlives the grace deadline, so that work cannot write after shutdown.
 
 ## Grace deadline enforcement
 

docs/actors/content/docs/lifecycle.mdx

@@ -347,21 +347,24 @@ When WebSocket hibernation is enabled, connections are preserved across sleep cy
 
 When an actor sleeps or is destroyed, it enters the graceful shutdown window:
 
-1. `c.abortSignal` fires and `c.aborted` becomes `true`. New connections and dispatch are rejected. Alarm timeouts are cancelled. On sleep, scheduled events are persisted and will be re-armed when the actor wakes.
-2. `onSleep` or `onDestroy` and `onDisconnect` for each closing connection run during the same window. User `waitUntil` promises and async raw WebSocket handlers are drained. Hibernatable WebSocket connections are preserved on sleep and closed on destroy.
-3. Once graceful work has completed, state is saved and final cleanup runs.
+1. `c.abortSignal` fires and `c.aborted` becomes `true`. New connections are rejected. Alarm timeouts are cancelled. On sleep, scheduled events are persisted and will be re-armed when the actor wakes.
+2. On sleep, actions, `keepAwake`, `waitUntil`, and WebSocket handlers that are already running finish before `onSleep` runs. Action results still reach the client. If this work does not finish within the first three quarters of `sleepGracePeriod`, `onSleep` runs anyway. On destroy, this step is skipped.
+3. `onSleep` or `onDestroy` and `onDisconnect` for each closing connection run during the same window. User `waitUntil` promises and async raw WebSocket handlers are drained. Hibernatable WebSocket connections are preserved on sleep and closed on destroy.
+4. Once graceful work has completed, state is saved and final cleanup runs.
 
 The entire window is bounded by `sleepGracePeriod` on both sleep and destroy. Defaults to 15 seconds. If the window is exceeded, the actor proceeds to state save anyway.
 
 #### Graceful shutdown window
 
-During steps 1 through 6, the actor is still fully functional. Database access, `broadcast`, `waitUntil`, `queue.send`, and `schedule.after` all work. State mutations are persisted at step 7. Actions invoked by pre-existing connections or lifecycle hooks continue to execute normally.
+During steps 1 through 3, the actor is still fully functional. Database access, `broadcast`, `waitUntil`, `queue.send`, and `schedule.after` all work. State mutations are persisted at step 4.
+
+On sleep, actions that arrive before `onSleep` runs can still run on the current instance. After that, actions sent through a handle are held and run on the next instance of the actor. An action sent through a connection while the connection closes fails with `ws.going_away`. It did not run, so it is safe to retry. On destroy, new actions are refused as soon as the shutdown sequence begins.
 
 New connections and raw WebSocket upgrades are rejected as soon as the shutdown sequence begins. New requests that arrive during shutdown are held until the actor wakes up again. The caller does not need to retry.
 
 If `schedule.after` is called during shutdown, the event is persisted so it survives the sleep/wake cycle, but no local timeout is scheduled. The event will fire after the actor wakes.
 
-In-flight actions are **not** waited on during shutdown. If an action must complete before the actor stops, wrap the critical work with `c.waitUntil()`.
+Raw `onRequest` handlers and queue sends do not delay `onSleep`. If that work must complete before the actor stops, wrap it with `c.waitUntil()`.
 
 ## Options
 

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -68,6 +68,10 @@ pub type ActionDispatchResult = std::result::Result<Vec<u8>, ActionDispatchError
 pub type HttpDispatchResult = Result<ActorHttpResponse>;
 
 const SERIALIZE_STATE_SHUTDOWN_SANITY_CAP: Duration = Duration::from_secs(15);
+/// Sleep grace keeps the last quarter of the grace period for onSleep, connection
+/// disconnects, and the final save, so in-flight work cannot use up the whole window.
+/// It scales with the configured grace period: the 15 second default leaves 3.75 seconds.
+const SLEEP_HOOK_RESERVE_DIVISOR: u32 = 4;
 #[cfg(test)]
 const LONG_SHUTDOWN_DRAIN_WARNING_THRESHOLD: Duration = Duration::from_secs(1);
 const INSPECTOR_SERIALIZE_STATE_INTERVAL: Duration = Duration::from_millis(50);
@@ -315,6 +319,18 @@ enum LiveExit {
 struct SleepGraceState {
 	deadline: Instant,
 	reason: ShutdownKind,
+	hooks: GraceHooks,
+}
+
+/// Whether onSleep and the connection disconnect hooks have run in this grace period.
+#[derive(Clone, Copy)]
+enum GraceHooks {
+	/// Sleep waits for work that started before grace, and runs the hooks at
+	/// `deadline` if that work is still running.
+	Pending {
+		deadline: Instant,
+	},
+	Emitted,
 }
 
 struct PersistedStartup {
@@ -486,6 +502,7 @@ impl ActorTask {
 			}
 			// TODO: Sample inbox depths periodically instead of on every loop iteration.
 			self.record_inbox_depths();
+			let hooks_deadline = self.pending_hooks_deadline();
 			tokio::select! {
 				biased;
 				lifecycle_command = self.lifecycle_inbox.recv() => {
@@ -522,6 +539,12 @@ impl ActorTask {
 						return exit;
 					}
 				}
+				_ = Self::sleep_grace_tick(hooks_deadline), if hooks_deadline.is_some() => {
+					self.emit_grace_hooks();
+					if let Some(exit) = self.try_finish_grace() {
+						return exit;
+					}
+				}
 				_ = Self::sleep_grace_tick(self.sleep_grace.as_ref().map(|grace| grace.deadline)), if self.sleep_grace.is_some() => {
 					if let Some(exit) = self.on_sleep_grace_deadline().await {
 						return exit;
@@ -672,8 +695,12 @@ impl ActorTask {
 			let activity = activity_notify.notified();
 			tokio::pin!(activity);
 
+			let hooks_deadline = self.pending_hooks_deadline();
 			tokio::select! {
 				_ = &mut activity => {}
+				_ = Self::sleep_grace_tick(hooks_deadline), if hooks_deadline.is_some() => {
+					self.emit_grace_hooks();
+				}
 				_ = Self::sleep_grace_tick(Some(deadline)) => {
 					if let Some(exit) = self.on_sleep_grace_deadline().await {
 						let LiveExit::Shutdown { reason } = exit else {
@@ -778,9 +805,39 @@ impl ActorTask {
 			ShutdownKind::Destroy => LifecycleState::DestroyGrace,
 		});
 		self.start_grace(reason);
+		match reason {
+			ShutdownKind::Sleep => self.try_emit_grace_hooks(),
+			ShutdownKind::Destroy => self.emit_grace_hooks(),
+		}
+	}
+
+	/// Runs the grace hooks once work that started before grace has finished.
+	/// The hooks deadline arm runs them otherwise.
+	fn try_emit_grace_hooks(&mut self) {
+		if self.ctx.shutdown_task_count() == 0 && self.ctx.websocket_callback_count() == 0 {
+			self.emit_grace_hooks();
+		}
+	}
+
+	fn emit_grace_hooks(&mut self) {
+		let Some(grace) = self.sleep_grace.as_mut() else {
+			return;
+		};
+		if matches!(grace.hooks, GraceHooks::Emitted) {
+			return;
+		}
+		grace.hooks = GraceHooks::Emitted;
+		let reason = grace.reason;
 		self.emit_grace_events(reason);
 	}
 
+	fn pending_hooks_deadline(&self) -> Option<Instant> {
+		match self.sleep_grace.as_ref()?.hooks {
+			GraceHooks::Pending { deadline } => Some(deadline),
+			GraceHooks::Emitted => None,
+		}
+	}
+
 	fn emit_grace_events(&mut self, reason: ShutdownKind) {
 		let conns: Vec<_> = self.ctx.conns().collect();
 		for conn in conns {
@@ -1159,10 +1216,10 @@ impl ActorTask {
 
 		match self.lifecycle {
 			LifecycleState::Started => None,
-			// The app's onSleep and connection cleanup still run, so the actor still serves work.
-			LifecycleState::SleepGrace if self.ctx.core_dispatched_hook_count() > 0 => None,
-			// The app has released its resources. The work never ran here, so the client
-			// retries it and the next generation serves it.
+			// Work admitted before onSleep runs is waited for before onSleep starts.
+			LifecycleState::SleepGrace if self.pending_hooks_deadline().is_some() => None,
+			// onSleep has started releasing the app's resources. The work never ran here,
+			// so the client retries it and the next generation serves it.
 			LifecycleState::SleepGrace | LifecycleState::SleepFinalize => {
 				self.ctx.warn_work_sent_to_stopping_instance("dispatch");
 				Some(ActorLifecycleError::Restarting.build())
@@ -1604,9 +1661,13 @@ impl ActorTask {
 		self.sleep_deadline = None;
 		self.ctx.cancel_sleep_timer();
 		self.ctx.cancel_actor_abort_signal();
+		let deadline = Instant::now() + grace_period;
 		self.sleep_grace = Some(SleepGraceState {
-			deadline: Instant::now() + grace_period,
+			deadline,
 			reason,
+			hooks: GraceHooks::Pending {
+				deadline: deadline - grace_period / SLEEP_HOOK_RESERVE_DIVISOR,
+			},
 		});
 		self.ctx.reset_sleep_timer();
 	}
@@ -1638,10 +1699,14 @@ impl ActorTask {
 	}
 
 	fn try_finish_grace(&mut self) -> Option<LiveExit> {
+		self.try_emit_grace_hooks();
 		let Some(grace) = self.sleep_grace.as_ref() else {
 			return None;
 		};
-		if self.ctx.can_finalize_shutdown(grace.reason) {
+		// Before the hooks run, the hook count is zero, so finalizing would skip onSleep.
+		if matches!(grace.hooks, GraceHooks::Emitted)
+			&& self.ctx.can_finalize_shutdown(grace.reason)
+		{
 			let reason = grace.reason;
 			self.sleep_grace = None;
 			return Some(LiveExit::Shutdown { reason });