From c2a8bff723eba254d57f2fd08d6d9d6909e5c6db Mon Sep 17 00:00:00 2001
From: Nathan Flurry <git@nathanflurry.com>
Date: Thu, 27 Aug 2026 15:58:15 -0700
Subject: [PATCH] fix(rivetkit): persist state across development restarts

---
 .../packages/rivetkit/src/context.rs          | 33 +++++++++++++++----
 rivetkit-rust/packages/rivetkit/src/start.rs  |  3 +-
 .../rivetkit/tests/modules/context.rs         | 16 ++++++++-
 3 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/rivetkit-rust/packages/rivetkit/src/context.rs b/rivetkit-rust/packages/rivetkit/src/context.rs
index 37235aa3b3..29c2eab05c 100644
--- a/rivetkit-rust/packages/rivetkit/src/context.rs
+++ b/rivetkit-rust/packages/rivetkit/src/context.rs
@@ -84,20 +84,33 @@ impl<S> Deref for StateRef<'_, S> {
 }
 
 pub struct StateMut<'a, S> {
-	guard: MappedRwLockWriteGuard<'a, S>,
+	guard: Option<MappedRwLockWriteGuard<'a, S>>,
+	inner: &'a ActorContext,
 }
 
 impl<S> Deref for StateMut<'_, S> {
 	type Target = S;
 
 	fn deref(&self) -> &Self::Target {
-		&self.guard
+		self.guard.as_deref().expect("state guard already dropped")
 	}
 }
 
 impl<S> DerefMut for StateMut<'_, S> {
 	fn deref_mut(&mut self) -> &mut Self::Target {
-		&mut self.guard
+		self.guard
+			.as_deref_mut()
+			.expect("state guard already dropped")
+	}
+}
+
+impl<S> Drop for StateMut<'_, S> {
+	fn drop(&mut self) {
+		// Release the state lock before scheduling serialization. This mirrors
+		// TypeScript's write-through state proxy and avoids relying on graceful
+		// process shutdown for persistence.
+		drop(self.guard.take());
+		self.inner.request_save(RequestSaveOpts::default());
 	}
 }
 
@@ -198,15 +211,22 @@ impl<A: Actor> Ctx<A> {
 	pub fn state_mut(&self) -> StateMut<'_, A::State> {
 		self.state.dirty.store(true, Ordering::Release);
 		StateMut {
-			guard: RwLockWriteGuard::map(self.state.value.write(), |state| {
+			guard: Some(RwLockWriteGuard::map(self.state.value.write(), |state| {
 				state.as_mut().expect("actor state not initialized")
-			}),
+			})),
+			inner: &self.inner,
 		}
 	}
 
 	pub fn set_state(&self, state: A::State) {
 		*self.state.value.write() = Some(state);
 		self.state.dirty.store(true, Ordering::Release);
+		self.inner.request_save(RequestSaveOpts::default());
+	}
+
+	pub(crate) fn set_initial_state(&self, state: A::State) {
+		*self.state.value.write() = Some(state);
+		self.clear_state_dirty();
 	}
 
 	pub fn state_dirty(&self) -> bool {
@@ -229,8 +249,7 @@ impl<A: Actor> Ctx<A> {
 	}
 
 	pub fn set_state_from_snapshot(&self, bytes: &[u8]) -> Result<()> {
-		self.set_state(Self::decode_state_snapshot(bytes)?);
-		self.clear_state_dirty();
+		self.set_initial_state(Self::decode_state_snapshot(bytes)?);
 		Ok(())
 	}
 
diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs
index 32e6c56d57..6de21aabff 100644
--- a/rivetkit-rust/packages/rivetkit/src/start.rs
+++ b/rivetkit-rust/packages/rivetkit/src/start.rs
@@ -205,8 +205,7 @@ pub async fn run_actor<A: Actor>(start: Start<A>) -> Result<()> {
 			// rivetkit-typescript where createState receives undefined input.
 			None => A::create_state(&ctx, input.decode_or_default()?).await?,
 		};
-		ctx.set_state(state);
-		ctx.clear_state_dirty();
+		ctx.set_initial_state(state);
 
 		let actor = Arc::new(A::create(&ctx).await?);
 		if is_new {
diff --git a/rivetkit-rust/packages/rivetkit/tests/modules/context.rs b/rivetkit-rust/packages/rivetkit/tests/modules/context.rs
index bbe2bde0d1..e46f96ba58 100644
--- a/rivetkit-rust/packages/rivetkit/tests/modules/context.rs
+++ b/rivetkit-rust/packages/rivetkit/tests/modules/context.rs
@@ -1,3 +1,7 @@
+use std::sync::{
+	Arc,
+	atomic::{AtomicUsize, Ordering},
+};
 use std::time::Duration;
 
 use serde::{Deserialize, Serialize};
@@ -76,8 +80,16 @@ fn typed_ctx_emit_accepts_named_events() {
 
 #[test]
 fn state_cell_reads_writes_and_tracks_dirty() {
+	let inner = actor_context("actor-id", "test", Vec::new(), "local");
+	let save_requests = Arc::new(AtomicUsize::new(0));
+	inner.on_request_save(Box::new({
+		let save_requests = Arc::clone(&save_requests);
+		move |_| {
+			save_requests.fetch_add(1, Ordering::SeqCst);
+		}
+	}));
 	let ctx = Ctx::<StatefulActor>::with_state(
-		actor_context("actor-id", "test", Vec::new(), "local"),
+		inner,
 		TestState {
 			count: 1,
 			label: "initial".into(),
@@ -94,6 +106,7 @@ fn state_cell_reads_writes_and_tracks_dirty() {
 	}
 
 	assert!(ctx.state_dirty());
+	assert_eq!(save_requests.load(Ordering::SeqCst), 1);
 	assert_eq!(
 		*ctx.state(),
 		TestState {
@@ -111,6 +124,7 @@ fn state_cell_reads_writes_and_tracks_dirty() {
 	});
 
 	assert!(ctx.state_dirty());
+	assert_eq!(save_requests.load(Ordering::SeqCst), 2);
 	assert_eq!(ctx.state().count, 7);
 }