fix(rivetkit): minor sleep-cleanup follow-ups

Author: NathanFlurry

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

@@ -439,10 +439,11 @@ impl ActorContext {
 			return Err(ActorLifecycleError::Stopping.build())
 				.context("destroy already requested for this generation");
 		}
-		self.cancel_sleep_timer();
-		self.flush_on_shutdown();
-		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.cancel();
+		// Reuse the shared teardown sequence used by the registry shutdown
+		// path so future changes to `mark_destroy_requested` cannot drift.
+		// `destroy_requested` is already true from the swap above; the redundant
+		// `store(true)` inside is harmless.
+		self.mark_destroy_requested();
 
 		let ctx = self.clone();
 		if Handle::try_current().is_ok() {

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

@@ -905,27 +905,65 @@ mod tests {
 		ctx.set_sleep_started(true);
 		ctx.set_ready(true);
 
+		// The stub must never flip the underlying flag.
 		ctx.set_prevent_sleep(true);
 		assert!(
 			!ctx.prevent_sleep(),
 			"prevent_sleep must stay false because the stub is a no-op"
 		);
-		// The sleep predicate ignores prevent_sleep entirely.
-		assert_eq!(ctx.can_sleep().await, CanSleep::Yes);
+
+		// Exhaustive match guards against reintroducing a `PreventSleep` enum
+		// variant. If a future change adds the variant back, this match stops
+		// compiling — surfacing the regression at build time rather than via a
+		// runtime assertion that could silently pass.
+		match ctx.can_sleep().await {
+			CanSleep::Yes
+			| CanSleep::NotReady
+			| CanSleep::NoSleep
+			| CanSleep::ActiveHttpRequests
+			| CanSleep::ActiveKeepAwake
+			| CanSleep::ActiveInternalKeepAwake
+			| CanSleep::ActiveRunHandler
+			| CanSleep::ActiveDisconnectCallbacks
+			| CanSleep::ActiveConnections
+			| CanSleep::ActiveWebSocketCallbacks => {}
+		}
 
 		ctx.set_prevent_sleep(false);
-		assert_eq!(ctx.can_sleep().await, CanSleep::Yes);
+		assert!(!ctx.prevent_sleep());
 	}
 
 	#[tokio::test(start_paused = true)]
-	async fn shutdown_deadline_token_cancels_on_request() {
+	async fn shutdown_deadline_token_aborts_select_awaiting_task() {
+		// Mirrors the NAPI `RunGracefulCleanup` pattern: a task awaits user
+		// work and the shutdown_deadline cancellation in a `tokio::select!`.
+		// If `cancel_shutdown_deadline()` does not propagate to clones of the
+		// token (a regression we cannot catch with `is_cancelled()` alone),
+		// the spawned task would hang and the test would time out.
 		let ctx = ActorContext::new_for_sleep_tests("actor-shutdown-deadline");
-
 		let token = ctx.shutdown_deadline_token();
 		assert!(!token.is_cancelled());
 
+		let aborted = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+		let aborted_clone = aborted.clone();
+		let task = tokio::spawn(async move {
+			tokio::select! {
+				_ = token.cancelled() => {
+					aborted_clone.store(true, Ordering::SeqCst);
+				}
+				_ = futures::future::pending::<()>() => {}
+			}
+		});
+
+		yield_now().await;
+		assert!(!aborted.load(Ordering::SeqCst));
+
 		ctx.cancel_shutdown_deadline();
-		assert!(token.is_cancelled());
+		task.await.expect("select task should join after cancel");
+		assert!(
+			aborted.load(Ordering::SeqCst),
+			"select-awaiting task must observe cancel via the cloned token"
+		);
 	}
 
 	#[tokio::test(start_paused = true)]

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

@@ -1671,7 +1671,13 @@ impl ActorTask {
 			return self.ctx.save_state(Vec::new()).await;
 		}
 
-		let deltas = match timeout(SERIALIZE_STATE_SHUTDOWN_SANITY_CAP, reply_rx).await {
+		// Cap at the larger of the default sanity bound or the user-configured
+		// sleep grace period. Without this, an actor with `sleepGracePeriod`
+		// raised above the default would silently truncate large state writes
+		// to empty deltas inside the user's own grace budget.
+		let cap = SERIALIZE_STATE_SHUTDOWN_SANITY_CAP
+			.max(self.factory.config().effective_sleep_grace_period());
+		let deltas = match timeout(cap, reply_rx).await {
 			Ok(Ok(Ok(deltas))) => deltas,
 			Ok(Ok(Err(error))) => {
 				tracing::error!(?error, "serializeState callback returned error");
@@ -1682,7 +1688,11 @@ impl ActorTask {
 				Vec::new()
 			}
 			Err(_) => {
-				tracing::error!("serializeState timed out");
+				tracing::error!(
+					actor_id = %self.ctx.actor_id(),
+					cap_ms = cap.as_millis() as u64,
+					"serializeState timed out; saving with empty deltas, prior persisted state retained"
+				);
 				Vec::new()
 			}
 		};

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

@@ -2678,8 +2678,11 @@ export class NativeActorContextAdapter {
 	}
 
 	destroy(): void {
-		markNativeDestroyRequested(this.#ctx);
+		// Call the native destroy first so it can throw `actor/starting` or
+		// `actor/stopping` without leaving an unresolved destroyCompletion
+		// promise behind in the native runtime state.
 		callNativeSync(() => this.#ctx.destroy());
+		markNativeDestroyRequested(this.#ctx);
 	}
 
 	client<T = AnyClient>(): T extends Registry<any> ? Client<T> : T {