Survive a widget/component slot flip inside one render call

reacton/core.py

@@ -1552,6 +1552,35 @@ def _possible_rerender(self):
         else:
             logger.info("No render phase triggered, already rendering")
 
+    def _discard_aborted_pass(self):
+        """Forget everything a render pass staged when it raised before reconciliation.
+
+        The *_next bookkeeping and the chained effects (previous_effect.next) of every
+        context belong to that pass. Mark every context dirty, so the next render()
+        re-executes the components and re-queues their effects from a consistent state.
+        Only for aborts in the render phase: a reconciliation that raised halfway already
+        created widgets, and those need the normal removal path.
+        """
+        # children_next is kept: it also holds contexts pre-created by state_set() (restored
+        # state), and a context created by the aborted pass owns no widgets yet, it is simply
+        # reused or pruned by the next render
+        contexts: List[ComponentContext] = [self.context_root]
+        while contexts:
+            context = contexts.pop()
+            for effect in context.effects:
+                effect.next = None
+            context.root_element_next = None
+            context.elements_next = {}
+            context.exceptions_self = []
+            context.exceptions_children = []
+            context.needs_render = True
+            context.needs_render_descendant = True
+            context.clean_subtree = False
+            contexts.extend(context.children.values())
+            contexts.extend(context.children_next.values())
+        self._shared_elements_next = set()
+        self.context = self.context_root
+
     def render(self, element: Element, container: widgets.Widget = None):
         # render + consolidate
         widget = None
@@ -1578,6 +1607,8 @@ def render(self, element: Element, container: widgets.Widget = None):
                 logger.info("Render requested on a closing/closed render context, ignoring")
                 return container
             prev_rc = getattr(local, "rc", None)
+            # an exception that escapes while this is True aborted a render pass (see the except below)
+            in_render_phase = True
             try:
                 local.rc = self
                 self.element = element
@@ -1653,10 +1684,12 @@ def format(reason: RerenderReason):
 
                         logger.info("Render reconsolidate...")
                         self.reconsolidating = True
+                        in_render_phase = False
                         try:
                             widget = self._reconsolidate(self.element, default_key="/", parent_key=ROOT_KEY)
                         finally:
                             self.reconsolidating = False
+                        in_render_phase = True
                         logger.info("Render reconsolidate done")
                         self.context.root_element = self.context.root_element_next
                         self.context.root_element_next = None
@@ -1705,7 +1738,16 @@ def format(reason: RerenderReason):
                     self._is_rendering = False
                 self.context = context_prev
                 logger.info("Done with render phase: %r", render_count)
-            except Exception as e:
+            except BaseException as e:
+                # Exceptions raised by components are collected in exceptions_self, so an
+                # exception here comes from the render machinery itself (duplicate key,
+                # hook misuse, ...) or is a cancellation, and aborted a pass halfway. Drop
+                # what that pass staged: otherwise the next render reconciles the last
+                # committed elements, but runs the effects the aborted pass chained, closed
+                # over elements that were never reconciled (get_widget then fails with
+                # "found in a previous render").
+                if in_render_phase:
+                    self._discard_aborted_pass()
                 if DEBUG:
                     # construct a fake traceback (showing how the elements were constructed)
                     if not self.tracebacks:
@@ -1845,8 +1887,13 @@ def _render(self, element: Element, default_key: str, parent_key: str):
             needs_render = context.needs_render
             if not needs_render:
                 if el_prev is not None and context_previous is context:
-                    assert not isinstance(el_prev.component, ComponentWidget)
-                    needs_render = el._arguments_changed(el_prev)
+                    if isinstance(el_prev.component, ComponentWidget):
+                        # an earlier pass of this render() call put a widget element at this
+                        # slot (elements_next), while the context is from the last reconciled
+                        # render: nothing to compare the arguments against, so render
+                        needs_render = True
+                    else:
+                        needs_render = el._arguments_changed(el_prev)
                 if context.exceptions_children:
                     # we have exceptions, so we need to render
                     needs_render = True
@@ -2505,8 +2552,11 @@ def _render(self, element: Element, default_key: str, parent_key: str):
 
         needs_render = context.needs_render
         if not needs_render and el_prev is not None and context_previous is context:
-            assert not isinstance(el_prev.component, ComponentWidget)
-            needs_render = el._arguments_changed(el_prev)
+            if isinstance(el_prev.component, ComponentWidget):
+                # see the classic renderer: a widget element from an earlier pass of this call
+                needs_render = True
+            else:
+                needs_render = el._arguments_changed(el_prev)
         if not needs_render and context.exceptions_children:
             # we have exceptions, so we need to render
             needs_render = True

reacton/core_test.py

@@ -3288,3 +3288,162 @@ def Test():
         set_state(1)
     rc.render(w.HTML(value="recover").key("HTML"))
     rc.close()
+
+
+def test_render_nested_pass_widget_to_component_flip():
+    # One render() call, two passes: the first pass puts a widget element at a
+    # slot, a state write during that pass opens a nested pass, and the nested
+    # pass puts back the function component that held the slot in the last
+    # committed render. _render then compares the component element against
+    # the widget element left in elements_next by the first pass and trips
+    # `assert not isinstance(el_prev.component, ComponentWidget)`.
+    set_mode = cast(Callable[[str], None], None)
+
+    set_clicks = cast(Callable[[int], None], None)
+
+    @react.component
+    def Child():
+        nonlocal set_clicks
+        clicks, set_clicks = react.use_state(0)
+        return w.Button(description=f"child {clicks}")
+
+    @react.component
+    def Parent():
+        nonlocal set_mode
+        mode, set_mode = react.use_state("component")
+        if mode == "widget":
+            # a set during render opens a nested pass in the same render() call
+            set_mode("component")
+            return w.Label(description="loading")
+        return Child()
+
+    box, rc = react.render(Parent(), handle_error=False)
+    set_clicks(3)
+    assert box.children[0].description == "child 3"
+    set_mode("widget")
+    assert isinstance(box.children[0], widgets.Button)
+    # nothing was committed in between, so Child was never unmounted and keeps its state
+    assert box.children[0].description == "child 3"
+    rc.close()
+
+
+def test_render_nested_pass_flip_strands_get_widget_effect():
+    # The production chain behind the flip above (Grotto PENG-1257, through
+    # solara's RoutingProvider): the AssertionError escapes render(), and a
+    # parent that re-executed in the aborted nested pass already chained a new
+    # effect closure (previous_effect.next) over its new root element, while
+    # its root_element_next was never set. The next render that does not
+    # re-execute the parent runs that closure, and get_widget on the never
+    # reconciled root raises "was found to be in a previous render".
+    set_mode = cast(Callable[[str], None], None)
+    set_counter = cast(Callable[[int], None], None)
+    set_other = cast(Callable[[int], None], None)
+
+    @react.component
+    def Child():
+        return w.Button(description="child")
+
+    @react.component
+    def Flipper():
+        nonlocal set_mode
+        mode, set_mode = react.use_state("component")
+        if mode == "widget":
+            # a set during render opens a nested pass, and the parent's state
+            # changes too, so the parent re-executes in that nested pass
+            # (in production: a router.push from a subscriber)
+            set_counter(1)
+            set_mode("component")
+            return w.Label(description="loading")
+        return Child()
+
+    @react.component
+    def Other():
+        nonlocal set_other
+        value, set_other = react.use_state(0)
+        return w.Label(description=f"other {value}")
+
+    @react.component
+    def Parent():
+        nonlocal set_counter
+        _counter, set_counter = react.use_state(0)
+
+        def effect():
+            react.get_widget(root)
+
+        react.use_effect(effect)
+        root = w.VBox(children=[Flipper(), Other()])
+        return root
+
+    box, rc = react.render(Parent(), handle_error=False)
+    root = box.children[0]
+    assert root.children[1].description == "other 0"
+    try:
+        set_mode("widget")
+    except AssertionError:
+        # the first stage, covered by test_render_nested_pass_widget_to_component_flip
+        pass
+    # an unrelated state change must not run a stale effect closure
+    set_other(1)
+    assert root.children[1].description == "other 1"
+    rc.close()
+
+
+@pytest.mark.parametrize("abort_in", ["first_pass", "nested_pass"])
+def test_render_escaped_exception_strands_chained_effect(abort_in):
+    # An exception raised by reacton's own bookkeeping inside _render (here a
+    # duplicate key) escapes render() and leaves the tree half updated: a
+    # component that re-executed in the aborted pass already chained a new
+    # effect closure (previous_effect.next) over the elements of that pass, but
+    # its root_element_next was never set. The next render that does not
+    # re-execute that component runs the chained effect against elements that
+    # were never reconciled; with get_widget that is the
+    # "was found to be in a previous render" KeyError.
+    # (Seen in production through solara's RoutingProvider, which calls
+    # get_widget(main) from an effect without dependencies.)
+    set_dup = cast(Callable[[bool], None], None)
+    set_other = cast(Callable[[int], None], None)
+    set_counter = cast(Callable[[int], None], None)
+
+    @react.component
+    def Other():
+        nonlocal set_other
+        value, set_other = react.use_state(0)
+        return w.Label(description=f"other {value}")
+
+    @react.component
+    def Dup(dup: bool, counter: int):
+        aborted = react.use_ref(False)
+        if dup and counter == 0 and abort_in == "nested_pass":
+            # make the parent re-execute in a nested pass of this render() call
+            set_counter(1)
+        abort_now = counter == (1 if abort_in == "nested_pass" else 0)
+        if dup and abort_now and not aborted.current:
+            # abort this pass with an exception from reacton's bookkeeping
+            aborted.current = True
+            return w.VBox(children=[w.Label(description="a").key("dup"), w.Label(description="b").key("dup")])
+        return w.Label(description="no dup")
+
+    @react.component
+    def Parent():
+        nonlocal set_dup, set_counter
+        dup, set_dup = react.use_state(False)
+        counter, set_counter = react.use_state(0)
+
+        def effect():
+            react.get_widget(root)
+
+        react.use_effect(effect)
+        root = w.VBox(children=[Other(), Dup(dup, counter)])
+        return root
+
+    box, rc = react.render(Parent(), handle_error=False)
+    root = box.children[0]
+    assert root.children[0].description == "other 0"
+    # pass 1: Dup bumps the counter, nested pass: Parent re-executes and chains
+    # a new effect closure over its new root, then Dup aborts the pass
+    with pytest.raises(KeyError, match="Duplicate key"):
+        set_dup(True)
+    # an unrelated state change must not run a stale effect closure
+    set_other(1)
+    assert root.children[0].description == "other 1"
+    rc.close()