state: subclass of MutableState must return _mark_dirty return value (#1898)

reflex/state.py

@@ -1868,7 +1868,13 @@ class ImmutableMutableProxy(MutableProxy):
     to modify the wrapped object when the StateProxy is immutable.
     """
 
-    def _mark_dirty(self, wrapped=None, instance=None, args=tuple(), kwargs=None):
+    def _mark_dirty(
+        self,
+        wrapped=None,
+        instance=None,
+        args=tuple(),
+        kwargs=None,
+    ) -> Any:
         """Raise an exception when an attempt is made to modify the object.
 
         Intended for use with `FunctionWrapper` from the `wrapt` library.
@@ -1879,6 +1885,9 @@ class ImmutableMutableProxy(MutableProxy):
             args: The args for the wrapped function.
             kwargs: The kwargs for the wrapped function.
 
+        Returns:
+            The result of the wrapped function.
+
         Raises:
             ImmutableStateError: if the StateProxy is not mutable.
         """
@@ -1887,6 +1896,6 @@ class ImmutableMutableProxy(MutableProxy):
                 "Background task StateProxy is immutable outside of a context "
                 "manager. Use `async with self` to modify state."
             )
-        super()._mark_dirty(
+        return super()._mark_dirty(
             wrapped=wrapped, instance=instance, args=args, kwargs=kwargs
         )

tests/test_state.py

@@ -1625,7 +1625,7 @@ class BackgroundTaskState(State):
     """A state with a background task."""
 
     order: List[str] = []
-    dict_list: Dict[str, List[int]] = {"foo": []}
+    dict_list: Dict[str, List[int]] = {"foo": [1, 2, 3]}
 
     @rx.background
     async def background_task(self):
@@ -1657,6 +1657,9 @@ class BackgroundTaskState(State):
                 pass  # update proxy instance
 
         async with self:
+            # Methods on ImmutableMutableProxy should return their wrapped return value.
+            assert self.dict_list.pop("foo") is not None
+
             self.order.append("background_task:stop")
             self.other()  # direct calling event handlers works in context
             self._private_method()