diff --git a/redwood-treehouse-guest/src/commonMain/kotlin/app/cash/redwood/treehouse/treehouseCompose.kt b/redwood-treehouse-guest/src/commonMain/kotlin/app/cash/redwood/treehouse/treehouseCompose.kt
index 0e9ec91a45..292ca50101 100644
--- a/redwood-treehouse-guest/src/commonMain/kotlin/app/cash/redwood/treehouse/treehouseCompose.kt
+++ b/redwood-treehouse-guest/src/commonMain/kotlin/app/cash/redwood/treehouse/treehouseCompose.kt
@@ -109,7 +109,7 @@ private class RedwoodZiplineTreehouseUi(
     )
   }
 
-  override fun snapshotState(): StateSnapshot? {
+  override fun snapshotState(): StateSnapshot {
     val savedState = saveableStateRegistry.performSave()
     return savedState.toStateSnapshot()
   }
diff --git a/redwood-treehouse/build.gradle b/redwood-treehouse/build.gradle
index b4b1f8bd43..5cb5e62f7a 100644
--- a/redwood-treehouse/build.gradle
+++ b/redwood-treehouse/build.gradle
@@ -21,6 +21,13 @@ kotlin {
         api libs.zipline
       }
     }
+
+    commonTest {
+      dependencies {
+        implementation libs.kotlin.test
+        implementation libs.assertk
+      }
+    }
   }
 }
 
diff --git a/redwood-treehouse/src/commonMain/kotlin/app/cash/redwood/treehouse/StateSnapshot.kt b/redwood-treehouse/src/commonMain/kotlin/app/cash/redwood/treehouse/StateSnapshot.kt
index a276588b6f..fa0b2393a2 100644
--- a/redwood-treehouse/src/commonMain/kotlin/app/cash/redwood/treehouse/StateSnapshot.kt
+++ b/redwood-treehouse/src/commonMain/kotlin/app/cash/redwood/treehouse/StateSnapshot.kt
@@ -18,20 +18,28 @@ package app.cash.redwood.treehouse
 import androidx.compose.runtime.MutableState
 import androidx.compose.runtime.mutableStateOf
 import kotlin.jvm.JvmInline
+import kotlinx.serialization.Contextual
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.JsonArray
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.JsonPrimitive
 import kotlinx.serialization.json.booleanOrNull
 import kotlinx.serialization.json.doubleOrNull
+import kotlinx.serialization.json.intOrNull
 
 @Serializable
 public class StateSnapshot(
-  public val content: Map<String, List<JsonElement?>>,
+  public val content: Map<String, List<@Contextual Saveable>>,
 ) {
-  public fun toValuesMap(): Map<String, List<Any?>>? {
+  public fun toValuesMap(): Map<String, List<Any?>> {
     return content.mapValues { entry ->
-      entry.value.map { mutableStateOf(it.fromJsonElement()) }
+      entry.value.map {
+        if (it.isMutableState) {
+          mutableStateOf(it.value)
+        } else {
+          it.value.fromJsonElement()
+        }
+      }
     }
   }
 
@@ -48,8 +56,8 @@ public fun Map<String, List<Any?>>.toStateSnapshot(): StateSnapshot = StateSnaps
   mapValues { entry ->
     entry.value.map { element ->
       when (element) {
-        is MutableState<*> -> element.value.toJsonElement()
-        else -> error("unexpected type: $this")
+        is MutableState<*> -> Saveable(true, element.value.toJsonElement())
+        else -> Saveable(false, element.toJsonElement())
       }
     }
   },
@@ -58,6 +66,7 @@ public fun Map<String, List<Any?>>.toStateSnapshot(): StateSnapshot = StateSnaps
 private fun Any?.toJsonElement(): JsonElement {
   return when (this) {
     is String -> JsonPrimitive(this)
+    is Int -> JsonPrimitive(this)
     is List<*> -> JsonArray(map { it.toJsonElement() })
     is JsonElement -> this
     else -> error("unexpected type: $this")
@@ -68,11 +77,9 @@ private fun Any?.toJsonElement(): JsonElement {
 private fun JsonElement?.fromJsonElement(): Any {
   return when (this) {
     is JsonPrimitive -> {
-      if (this.isString) {
-        return content
-      }
-      return booleanOrNull ?: doubleOrNull ?: error("unexpected type: $this")
-      // TODO add other primitive types (double, float, long) when needed
+      if (this.isString) return content
+      return booleanOrNull ?: doubleOrNull ?: intOrNull ?: error("unexpected type: $this")
+      // TODO add other primitive types (float, long) when needed
     }
 
     is JsonArray -> listOf({ this.forEach { it.toJsonElement() } })
@@ -81,3 +88,8 @@ private fun JsonElement?.fromJsonElement(): Any {
     else -> error("unexpected type: $this")
   }
 }
+
+public data class Saveable(
+  val isMutableState: Boolean,
+  val value: JsonElement,
+)
diff --git a/redwood-treehouse/src/commonTest/kotlin/app/cash/redwood/treehouse/StateSnapshotTest.kt b/redwood-treehouse/src/commonTest/kotlin/app/cash/redwood/treehouse/StateSnapshotTest.kt
new file mode 100644
index 0000000000..c25c93f534
--- /dev/null
+++ b/redwood-treehouse/src/commonTest/kotlin/app/cash/redwood/treehouse/StateSnapshotTest.kt
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2023 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package app.cash.redwood.treehouse
+
+import androidx.compose.runtime.MutableState
+import androidx.compose.runtime.mutableStateOf
+import assertk.assertThat
+import assertk.assertions.containsOnly
+import assertk.assertions.isEqualTo
+import kotlin.test.Test
+import kotlin.test.assertTrue
+import kotlinx.serialization.json.JsonPrimitive
+
+class StateSnapshotTest {
+
+  @Test
+  fun toValueMapWorksAsExpected() {
+    val stateSnapshot = stateSnapshot()
+    val valuesMap = stateSnapshot.toValuesMap()
+    assertThat(valuesMap.entries.size).isEqualTo(4)
+    assertTrue(valuesMap["key1"]!![0] is MutableState<*>)
+    assertThat((valuesMap["key1"]!![0] as MutableState<*>).value).isEqualTo(JsonPrimitive(1))
+
+    assertThat(valuesMap["key2"]).isEqualTo(listOf(1.0))
+
+    assertThat(valuesMap["key3"]!![0] is MutableState<*>)
+    assertThat((valuesMap["key3"]!![0] as MutableState<*>).value).isEqualTo(JsonPrimitive("str"))
+
+    assertThat(valuesMap["key4"]).isEqualTo(listOf("str"))
+  }
+
+  @Test
+  fun toStateSnapshotWorksAsExpected() {
+    val storedStateSnapshot = storedStateSnapshot()
+    val stateSnapshot = storedStateSnapshot.toStateSnapshot()
+    assertThat(stateSnapshot.content).containsOnly(
+      "key1" to listOf(Saveable(true, JsonPrimitive(1))),
+      "key2" to listOf(Saveable(false, JsonPrimitive(1))),
+      "key3" to listOf(Saveable(true, JsonPrimitive("str"))),
+      "key4" to listOf(Saveable(false, JsonPrimitive("str"))),
+    )
+  }
+
+  private fun stateSnapshot() = StateSnapshot(
+    mapOf(
+      "key1" to listOf(Saveable(true, JsonPrimitive(1))),
+      "key2" to listOf(Saveable(false, JsonPrimitive(1))),
+      "key3" to listOf(Saveable(true, JsonPrimitive("str"))),
+      "key4" to listOf(Saveable(false, JsonPrimitive("str"))),
+    ),
+  )
+
+  private fun storedStateSnapshot() = mapOf(
+    "key1" to listOf(mutableStateOf(1)),
+    "key2" to listOf(1),
+    "key3" to listOf(mutableStateOf("str")),
+    "key4" to listOf("str"),
+  )
+}