Fix branching off an async rx node

param/reactive.py

@@ -1615,7 +1615,8 @@ def __init__(
         self._trigger: Trigger | None
         if operation and (iscoroutinefunction(operation['fn']) or inspect.isgeneratorfunction(operation['fn'])):
             self._trigger = Trigger(internal=True)
-            self._current_ = Undefined
+            self._current_ = Undefined
+            self._dirty = True  # Otherwise current will be stuck as Undefined.
         else:
             self._trigger = None
         self._root = self._compute_root()
@@ -1907,15 +1908,25 @@ def _resolve(self):
                     # If this rx is cloned from an shared input then we make use
                     # of the shared.rx.value to ensure branching pipelines do
                     # not have to recompute the inputs multiple times.
-                    if self._is_async:
-                        self._shared.rx.value # trigger async resolve
+                    shared = self._shared
+                    value = shared.rx.value # triggers async resolve
+                    if self._is_async and (
+                        shared._awaiting or shared._current_task is not None
+                    ):
+                        # The shared node is still processing, resolve when finished
                         self._lazy_resolve()
                         raise Skip
                     # Returns instead of raising Skip because this path does
                     # resolve to a value, so it must mirror the shared node's
                     # skip state rather than be marked skipped by the handler.
-                    self._current_ = self._shared.rx.value
-                    self._skipped = self._shared._skipped
+                    self._current_ = value
+                    self._skipped = shared._skipped
+                    if self._is_async:
+                        # The value was adopted without scheduling a task, so
+                        # claim a generation for it. This supersedes a task an
+                        # earlier operation may have scheduled and still awaits a resolution.
+                        self._resolve_generation += 1
+                        self._finished_generation = self._resolve_generation
                     self._dirty = False
                     return self._current_
                 operation = self._operation

tests/testreactive.py

@@ -1069,6 +1069,100 @@ async def expensive_compute(a):
     # the superseded updates are not computed at all.
     assert call_count == 2
 
+async def _pair(value):
+    await asyncio.sleep(0.02)
+    return (value * 2, value * 3)
+
+async def test_async_shared_rx_branch_after_settling_resolves():
+    irx = rx(1)
+    node = irx.rx.pipe(_pair)
+    node.rx.value
+    await async_wait_until(lambda: node.rx.value == (2, 3))
+
+    # The branch mirrors a node that has already settled, so it must resolve
+    # on its first read rather than being stranded on Undefined.
+    first = node[0]
+    assert first.rx.value == 2
+
+    irx.rx.value = 3
+    await async_wait_until(lambda: first.rx.value == 6)
+
+async def test_async_shared_rx_branch_while_awaiting_resolves():
+    irx = rx(1)
+    node = irx.rx.pipe(_pair)
+    node.rx.value
+
+    first = node[0]
+    assert first.rx.value is param.Undefined
+    await async_wait_until(lambda: first.rx.value == 2)
+
+async def test_async_shared_rx_branch_before_resolving_resolves():
+    irx = rx(1)
+    node = irx.rx.pipe(_pair)
+    first, second = node[0], node[1]
+
+    assert first.rx.value is param.Undefined
+    await async_wait_until(lambda: first.rx.value == 2)
+
+    # The shared node has settled by now, so the sibling branch resolves on
+    # its first read.
+    assert second.rx.value == 3
+
+async def test_async_shared_rx_branch_computed_once():
+    call_count = 0
+
+    async def count_pair(value):
+        nonlocal call_count
+        call_count += 1
+        await asyncio.sleep(0.02)
+        return (value * 2, value * 3)
+
+    irx = rx(1)
+    node = irx.rx.pipe(count_pair)
+    first, second = node[0], node[1]
+
+    # Request the value for both nodes
+    first.rx.value
+    second.rx.value
+    await async_wait_until(lambda: first.rx.value == 2 and second.rx.value == 3)
+
+    # The branches resolve through the shared node rather than recomputing.
+    assert call_count == 1
+
+    irx.rx.value = 3
+    await async_wait_until(lambda: first.rx.value == 6 and second.rx.value == 9)
+    assert call_count == 2
+
+async def test_async_shared_rx_branch_notifies_watcher():
+    irx = rx(1)
+    node = irx.rx.pipe(_pair)
+    node.rx.value
+    await async_wait_until(lambda: node.rx.value == (2, 3))
+
+    items = []
+    first = node[0]
+    assert first.rx.value == 2
+    first.rx.watch(items.append)
+
+    irx.rx.value = 3
+    await async_wait_until(lambda: items == [6])
+    assert items == [6]
+
+async def test_async_gen_shared_rx_branch_resolves():
+    async def gen(value):
+        for i in range(3):
+            await asyncio.sleep(0.02)
+            yield (value + i, i)
+
+    irx = rx(1)
+    node = irx.rx.pipe(gen)
+    node.rx.value
+    await async_wait_until(lambda: node.rx.value == (3, 2))
+
+    # A branch of a generator node adopts the value the generator settled on.
+    first = node[0]
+    assert first.rx.value == 3
+
 @pytest.mark.parametrize('lazy', [False, True])
 def test_root_invalidation(lazy):
     arx = rx('a', lazy=lazy)