Repository: reflex-dev/reflex

Issue:
[REF-3184] All my rx.background don't work on 0.5.5
Hello everyone

All my rx.background don't work on 0.5.5

You can try make any foreach with rx.background and yield Anystate.do_this and it won't work

I can show my code to DM any contributors 

<sub>[REF-3184](https://linear.app/reflex-dev/issue/REF-3184/all-my-rxbackground-dont-work-on-055)</sub>


Discussion:
@fungalspores can you share your code with me, `masen@reflex.dev`. Or share a minimal reproducer here in this issue.

I'm not able to reproduce the issue with the following code

```python
import asyncio
import reflex as rx


class MyTaskState(rx.State):
    counter: int = 0
    max_counter: int = 10
    running: bool = False
    _n_tasks: int = 0

    def foo(self):
        print("foo")

    @rx.background
    async def my_task(self):
        async with self:
            # The latest state values are always available inside the context
            if self._n_tasks > 0:
                # only allow 1 concurrent task
                return

            # State mutation is only allowed inside context block
            self._n_tasks += 1

        while True:
            async with self:
                # Check for stopping conditions inside context
                if self.counter >= self.max_counter:
                    self.running = False
                if not self.running:
                    self._n_tasks -= 1
                    return

                self.counter += 1
                yield MyTaskState.foo

            # Await long operations outside the context to avoid blocking UI
            await asyncio.sleep(0.5)

    def toggle_running(self):
        self.running = not self.running
        if self.running:
            return MyTaskState.my_task

    def clear_counter(self):
        self.counter = 0


def background_task_example():
    return rx.hstack(
        rx.heading(MyTaskState.counter, " /"),
        rx.chakra.number_input(
            value=MyTaskState.max_counter,
            on_change=MyTaskState.set_max_counter,
            width="8em",
        ),
        rx.button(
            rx.cond(~MyTaskState.running, "Start", "Stop"),
            on_click=MyTaskState.toggle_running,
        ),
        rx.button(
            "Reset",
            on_click=MyTaskState.clear_counter,
        ),
    )
```

I think I'm missing the part about the `foreach`, not sure how that looks for the issue you are hitting.
about `return MyTaskState.my_task` code, I have another question:
In my memory, function `toggle_running` return `MyTaskState.my_task`, will send a message to web page, and then web page's applyEvent send back to server, and server run `my_task` in `_process_background`. 
my question: can it be more direct no through the web? or is there a function maybe name `run_tasks` can use to direct run a backgroup task? @masenf 
> @fungalspores can you share your code with me, `masen@reflex.dev`. Or share a minimal reproducer here in this issue.
> 
> I'm not able to reproduce the issue with the following code
> 
> ```python
> import asyncio
> import reflex as rx
> 
> 
> class MyTaskState(rx.State):
>     counter: int = 0
>     max_counter: int = 10
>     running: bool = False
>     _n_tasks: int = 0
> 
>     def foo(self):
>         print("foo")
> 
>     @rx.background
>     async def my_task(self):
>         async with self:
>             # The latest state values are always available inside the context
>             if self._n_tasks > 0:
>                 # only allow 1 concurrent task
>                 return
> 
>             # State mutation is only allowed inside context block
>             self._n_tasks += 1
> 
>         while True:
>             async with self:
>                 # Check for stopping conditions inside context
>                 if self.counter >= self.max_counter:
>                     self.running = False
>                 if not self.running:
>                     self._n_tasks -= 1
>                     return
> 
>                 self.counter += 1
>                 yield MyTaskState.foo
> 
>             # Await long operations outside the context to avoid blocking UI
>             await asyncio.sleep(0.5)
> 
>     def toggle_running(self):
>         self.running = not self.running
>         if self.running:
>             return MyTaskState.my_task
> 
>     def clear_counter(self):
>         self.counter = 0
> 
> 
> def background_task_example():
>     return rx.hstack(
>         rx.heading(MyTaskState.counter, " /"),
>         rx.chakra.number_input(
>             value=MyTaskState.max_counter,
>             on_change=MyTaskState.set_max_counter,
>             width="8em",
>         ),
>         rx.button(
>             rx.cond(~MyTaskState.running, "Start", "Stop"),
>             on_click=MyTaskState.toggle_running,
>         ),
>         rx.button(
>             "Reset",
>             on_click=MyTaskState.clear_counter,
>         ),
>     )
> ```
> 
> I think I'm missing the part about the `foreach`, not sure how that looks for the issue you are hitting.

I wrote to your mail and try added to discord friend for faster communication 
I had a similar issue: I was able to resolve it by removing the nested ` async with self:`

```
 @rx.background
 async def fetch_data(self):
        async with self:
            p = await self.get_state(PState)
            self.p_id = p.get_current_p_id 
            if self.p_id == "":
                return rx.redirect("/ws")
            wf = await self.fetch_flow()
            self.wf_id = wf
            while self.wf_status != "Completed":
                wf_status = await self.fetch_flow_run_id()
                async with self:
                    self.wf_status = wf_status["state"]
....
```
Previously I had an issue if the state not updating in my while -> so I added the extra `async with self:` in the while loop. Now with the new update the state did not want to update because of the extra `async with self:` so I removed it.
In my case if you add print statements you will see everything updating up to the second async with self and then not updating after it. 
@Ernelene thank you for the insight!

It appears this issue was introduced by https://github.com/reflex-dev/reflex/pull/3508 where i added locking around the `async with self` to avoid concurrent access



Patch:
diff --git a/reflex/state.py b/reflex/state.py
index 49b5bd4a413..ceadb621004 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -201,7 +201,7 @@ async def _no_chain_background_task_gen(*args, **kwargs):
 
 def _substate_key(
     token: str,
-    state_cls_or_name: BaseState | Type[BaseState] | str | list[str],
+    state_cls_or_name: BaseState | Type[BaseState] | str | Sequence[str],
 ) -> str:
     """Get the substate key.
 
@@ -2028,19 +2028,38 @@ async def bg_increment(self):
                     self.counter += 1
     """
 
-    def __init__(self, state_instance):
+    def __init__(
+        self, state_instance, parent_state_proxy: Optional["StateProxy"] = None
+    ):
         """Create a proxy for a state instance.
 
+        If `get_state` is used on a StateProxy, the resulting state will be
+        linked to the given state via parent_state_proxy. The first state in the
+        chain is the state that initiated the background task.
+
         Args:
             state_instance: The state instance to proxy.
+            parent_state_proxy: The parent state proxy, for linked mutability and context tracking.
         """
         super().__init__(state_instance)
         # compile is not relevant to backend logic
         self._self_app = getattr(prerequisites.get_app(), constants.CompileVars.APP)
-        self._self_substate_path = state_instance.get_full_name().split(".")
+        self._self_substate_path = tuple(state_instance.get_full_name().split("."))
         self._self_actx = None
         self._self_mutable = False
         self._self_actx_lock = asyncio.Lock()
+        self._self_actx_lock_holder = None
+        self._self_parent_state_proxy = parent_state_proxy
+
+    def _is_mutable(self) -> bool:
+        """Check if the state is mutable.
+
+        Returns:
+            Whether the state is mutable.
+        """
+        if self._self_parent_state_proxy is not None:
+            return self._self_parent_state_proxy._is_mutable()
+        return self._self_mutable
 
     async def __aenter__(self) -> StateProxy:
         """Enter the async context manager protocol.
@@ -2053,8 +2072,31 @@ async def __aenter__(self) -> StateProxy:
 
         Returns:
             This StateProxy instance in mutable mode.
-        """
+
+        Raises:
+            ImmutableStateError: If the state is already mutable.
+        """
+        if self._self_parent_state_proxy is not None:
+            parent_state = (
+                await self._self_parent_state_proxy.__aenter__()
+            ).__wrapped__
+            super().__setattr__(
+                "__wrapped__",
+                await parent_state.get_state(
+                    State.get_class_substate(self._self_substate_path)
+                ),
+            )
+            return self
+        current_task = asyncio.current_task()
+        if (
+            self._self_actx_lock.locked()
+            and current_task == self._self_actx_lock_holder
+        ):
+            raise ImmutableStateError(
+                "The state is already mutable. Do not nest `async with self` blocks."
+            )
         await self._self_actx_lock.acquire()
+        self._self_actx_lock_holder = current_task
         self._self_actx = self._self_app.modify_state(
             token=_substate_key(
                 self.__wrapped__.router.session.client_token,
@@ -2076,12 +2118,16 @@ async def __aexit__(self, *exc_info: Any) -> None:
         Args:
             exc_info: The exception info tuple.
         """
+        if self._self_parent_state_proxy is not None:
+            await self._self_parent_state_proxy.__aexit__(*exc_info)
+            return
         if self._self_actx is None:
             return
         self._self_mutable = False
         try:
             await self._self_actx.__aexit__(*exc_info)
         finally:
+            self._self_actx_lock_holder = None
             self._self_actx_lock.release()
         self._self_actx = None
 
@@ -2116,7 +2162,7 @@ def __getattr__(self, name: str) -> Any:
         Raises:
             ImmutableStateError: If the state is not in mutable mode.
         """
-        if name in ["substates", "parent_state"] and not self._self_mutable:
+        if name in ["substates", "parent_state"] and not self._is_mutable():
             raise ImmutableStateError(
                 "Background task StateProxy is immutable outside of a context "
                 "manager. Use `async with self` to modify state."
@@ -2156,7 +2202,7 @@ def __setattr__(self, name: str, value: Any) -> None:
         """
         if (
             name.startswith("_self_")  # wrapper attribute
-            or self._self_mutable  # lock held
+            or self._is_mutable()  # lock held
             # non-persisted state attribute
             or name in self.__wrapped__.get_skip_vars()
         ):
@@ -2180,7 +2226,7 @@ def get_substate(self, path: Sequence[str]) -> BaseState:
         Raises:
             ImmutableStateError: If the state is not in mutable mode.
         """
-        if not self._self_mutable:
+        if not self._is_mutable():
             raise ImmutableStateError(
                 "Background task StateProxy is immutable outside of a context "
                 "manager. Use `async with self` to modify state."
@@ -2199,12 +2245,14 @@ async def get_state(self, state_cls: Type[BaseState]) -> BaseState:
         Raises:
             ImmutableStateError: If the state is not in mutable mode.
         """
-        if not self._self_mutable:
+        if not self._is_mutable():
             raise ImmutableStateError(
                 "Background task StateProxy is immutable outside of a context "
                 "manager. Use `async with self` to modify state."
             )
-        return await self.__wrapped__.get_state(state_cls)
+        return type(self)(
+            await self.__wrapped__.get_state(state_cls), parent_state_proxy=self
+        )
 
     def _as_state_update(self, *args, **kwargs) -> StateUpdate:
         """Temporarily allow mutability to access parent_state.


Test patch:
diff --git a/integration/test_background_task.py b/integration/test_background_task.py
index 96a47e951fc..98b6e48ffeb 100644
--- a/integration/test_background_task.py
+++ b/integration/test_background_task.py
@@ -12,7 +12,10 @@ def BackgroundTask():
     """Test that background tasks work as expected."""
     import asyncio
 
+    import pytest
+
     import reflex as rx
+    from reflex.state import ImmutableStateError
 
     class State(rx.State):
         counter: int = 0
@@ -71,6 +74,38 @@ async def handle_racy_event(self):
                 self.racy_task(), self.racy_task(), self.racy_task(), self.racy_task()
             )
 
+        @rx.background
+        async def nested_async_with_self(self):
+            async with self:
+                self.counter += 1
+                with pytest.raises(ImmutableStateError):
+                    async with self:
+                        self.counter += 1
+
+        async def triple_count(self):
+            third_state = await self.get_state(ThirdState)
+            await third_state._triple_count()
+
+    class OtherState(rx.State):
+        @rx.background
+        async def get_other_state(self):
+            async with self:
+                state = await self.get_state(State)
+                state.counter += 1
+                await state.triple_count()
+            with pytest.raises(ImmutableStateError):
+                await state.triple_count()
+            with pytest.raises(ImmutableStateError):
+                state.counter += 1
+            async with state:
+                state.counter += 1
+                await state.triple_count()
+
+    class ThirdState(rx.State):
+        async def _triple_count(self):
+            state = await self.get_state(State)
+            state.counter *= 3
+
     def index() -> rx.Component:
         return rx.vstack(
             rx.chakra.input(
@@ -109,6 +144,16 @@ def index() -> rx.Component:
                 on_click=State.handle_racy_event,
                 id="racy-increment",
             ),
+            rx.button(
+                "Nested Async with Self",
+                on_click=State.nested_async_with_self,
+                id="nested-async-with-self",
+            ),
+            rx.button(
+                "Increment from OtherState",
+                on_click=OtherState.get_other_state,
+                id="increment-from-other-state",
+            ),
             rx.button("Reset", on_click=State.reset_counter, id="reset"),
         )
 
@@ -230,3 +275,61 @@ def test_background_task(
     assert background_task._poll_for(
         lambda: not background_task.app_instance.background_tasks  # type: ignore
     )
+
+
+def test_nested_async_with_self(
+    background_task: AppHarness,
+    driver: WebDriver,
+    token: str,
+):
+    """Test that nested async with self in the same coroutine raises Exception.
+
+    Args:
+        background_task: harness for BackgroundTask app.
+        driver: WebDriver instance.
+        token: The token for the connected client.
+    """
+    assert background_task.app_instance is not None
+
+    # get a reference to all buttons
+    nested_async_with_self_button = driver.find_element(By.ID, "nested-async-with-self")
+    increment_button = driver.find_element(By.ID, "increment")
+
+    # get a reference to the counter
+    counter = driver.find_element(By.ID, "counter")
+    assert background_task._poll_for(lambda: counter.text == "0", timeout=5)
+
+    nested_async_with_self_button.click()
+    assert background_task._poll_for(lambda: counter.text == "1", timeout=5)
+
+    increment_button.click()
+    assert background_task._poll_for(lambda: counter.text == "2", timeout=5)
+
+
+def test_get_state(
+    background_task: AppHarness,
+    driver: WebDriver,
+    token: str,
+):
+    """Test that get_state returns a state bound to the correct StateProxy.
+
+    Args:
+        background_task: harness for BackgroundTask app.
+        driver: WebDriver instance.
+        token: The token for the connected client.
+    """
+    assert background_task.app_instance is not None
+
+    # get a reference to all buttons
+    other_state_button = driver.find_element(By.ID, "increment-from-other-state")
+    increment_button = driver.find_element(By.ID, "increment")
+
+    # get a reference to the counter
+    counter = driver.find_element(By.ID, "counter")
+    assert background_task._poll_for(lambda: counter.text == "0", timeout=5)
+
+    other_state_button.click()
+    assert background_task._poll_for(lambda: counter.text == "12", timeout=5)
+
+    increment_button.click()
+    assert background_task._poll_for(lambda: counter.text == "13", timeout=5)
diff --git a/tests/test_state.py b/tests/test_state.py
index 2fc14938918..95bd4a93e2d 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -1806,7 +1806,7 @@ async def test_state_proxy(grandchild_state: GrandchildState, mock_app: rx.App):
 
     sp = StateProxy(grandchild_state)
     assert sp.__wrapped__ == grandchild_state
-    assert sp._self_substate_path == grandchild_state.get_full_name().split(".")
+    assert sp._self_substate_path == tuple(grandchild_state.get_full_name().split("."))
     assert sp._self_app is mock_app
     assert not sp._self_mutable
     assert sp._self_actx is None