fix: restore kernel context on reactive event handler

If a reactive variable was set from a different (or no) kernel context,
the event handler would not restore the kernel context that was
active when the event handler was registered.

Fixes #961

Author: maartenbreddels

solara/toestand.py

@@ -1,4 +1,3 @@
-import contextlib
 import dataclasses
 import inspect
 import logging
@@ -33,6 +32,7 @@
 import solara
 import solara.settings
 from solara import _using_solara_server
+from solara.util import nullcontext
 
 T = TypeVar("T")
 TS = TypeVar("TS")
@@ -176,18 +176,18 @@ def fire(self, new: T, old: T):
         scope_id = self._get_scope_key()
         scopes = set()
         for listener, scope in self.listeners[scope_id].copy():
-            if scope is not None:
-                scopes.add(scope)
+            scopes.add(scope)
         for listener2, scope in self.listeners2[scope_id].copy():
-            if scope is not None:
-                scopes.add(scope)
-        with contextlib.ExitStack() as stack:
+            scopes.add(scope)
+        if scopes:
             for scope in scopes:
-                stack.enter_context(scope)
-            for listener, scope in self.listeners[scope_id].copy():
-                listener(new)
-            for listener2, scope in self.listeners2[scope_id].copy():
-                listener2(new, old)
+                with scope or nullcontext():
+                    for listener, scope_listener in self.listeners[scope_id].copy():
+                        if scope == scope_listener:
+                            listener(new)
+                    for listener2, scope_listener in self.listeners2[scope_id].copy():
+                        if scope == scope_listener:
+                            listener2(new, old)
 
     def update(self, _f=None, **kwargs):
         if _f is not None:
@@ -885,6 +885,35 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         thread_local.reactive_used = self.reactive_used_before
 
 
+class Context:
+    def __init__(self, render_context, kernel_context):
+        # combine the render context *and* the kernel context into one context
+        self.render_context = render_context
+        self.kernel_context = kernel_context
+
+    def __enter__(self):
+        self.render_context.__enter__()
+        self.kernel_context.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # this will trigger a render
+        res1 = self.render_context.__exit__(exc_type, exc_val, exc_tb)
+        # pop the current context from the stack
+        res2 = self.kernel_context.__exit__(exc_type, exc_val, exc_tb)
+        return res1 or res2
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, Context):
+            return False
+        return self.render_context == value.render_context and self.kernel_context == value.kernel_context
+
+    def __hash__(self) -> int:
+        return hash(id(self.render_context)) ^ hash(id(self.kernel_context))
+
+    def __repr__(self) -> str:
+        return f"Context(render_context={self.render_context}, kernel_context={self.kernel_context})"
+
+
 class AutoSubscribeContextManagerReacton(AutoSubscribeContextManagerBase):
     def __init__(self, element: solara.Element):
         self.element = element
@@ -900,7 +929,9 @@ def force_update(new_value, old_value):
         super().__enter__()
 
         def update_subscribers():
-            self.update_subscribers(force_update, scope=reacton.core.get_render_context(required=True))
+            rc = reacton.core.get_render_context(required=True)
+            kernel = solara.server.kernel_context.get_current_context() if solara.server.kernel_context.has_current_context() else nullcontext()
+            self.update_subscribers(force_update, scope=Context(rc, kernel))
 
         solara.use_effect(update_subscribers, None)
 

tests/unit/toestand_test.py

@@ -15,6 +15,7 @@
 import solara as sol
 import solara.lab
 import solara.toestand as toestand
+import solara.server.kernel_context
 from solara.server import kernel, kernel_context
 from solara.toestand import Reactive, Ref, State, use_sync_external_store
 import solara.settings
@@ -180,25 +181,30 @@ def test_scopes(no_kernel_context):
     @solara.component
     def Page():
         reactive = solara.use_reactive(0)
+        assert solara.server.kernel_context.get_current_context() is expected_context
         if reactives[index] is None:
             reactives[index] = reactive  # type: ignore
         solara.Info(str(reactive.value))
 
     with context1:
+        expected_context = context1
         _, rc1 = react.render(Page(), handle_error=False)
         assert rc1.find(v.Alert).widget.children[0] == "0"
         assert reactives[0] is not None
 
     with context2:
         index = 1
+        expected_context = context2
         _, rc2 = react.render(Page(), handle_error=False)
         assert rc2.find(v.Alert).widget.children[0] == "0"
         assert reactives[1] is not None
 
+    expected_context = context1
     reactives[0].value = 1
     assert rc1.find(v.Alert).widget.children[0] == "1"
     assert rc2.find(v.Alert).widget.children[0] == "0"
 
+    expected_context = context2
     reactives[1].value = 2
     assert rc1.find(v.Alert).widget.children[0] == "1"
     assert rc2.find(v.Alert).widget.children[0] == "2"
@@ -215,6 +221,52 @@ def Page():
         rc2.close()
 
 
+def test_scopes_restore(no_kernel_context):
+    kernel1 = kernel.Kernel()
+    kernel2 = kernel.Kernel()
+    assert kernel_context.current_context[kernel_context.get_current_thread_key()] is None
+
+    context1 = kernel_context.VirtualKernelContext(id="toestand-1", kernel=kernel1, session_id="session-1")
+    context2 = kernel_context.VirtualKernelContext(id="toestand-2", kernel=kernel2, session_id="session-2")
+
+    from solara._stores import SharedStore
+
+    reactive: solara.Reactive[int] = solara.Reactive(SharedStore(10))
+
+    @solara.component
+    def Page1():
+        assert solara.server.kernel_context.get_current_context() is context1
+        solara.Info(str(reactive.value))
+        Page1b()
+
+    @solara.component
+    def Page1b():
+        assert solara.server.kernel_context.get_current_context() is context1
+        solara.Error(str(reactive.value))
+
+    @solara.component
+    def Page2():
+        assert solara.server.kernel_context.get_current_context() is context2
+        solara.Info(str(reactive.value))
+
+    with context2:
+        _, rc2 = react.render(Page2(), handle_error=False)
+        assert rc2.find(v.Alert).widget.children[0] == "10"
+
+    with context1:
+        _, rc1 = react.render(Page1(), handle_error=False)
+        assert rc1.find(v.Alert)[0].widget.children[0] == "10"
+
+    reactive.value = 42
+    assert rc1.find(v.Alert)[0].widget.children[0] == "42"
+    assert rc2.find(v.Alert).widget.children[0] == "42"
+
+    with context1:
+        rc1.close()
+    with context2:
+        rc2.close()
+
+
 def test_nested_update():
     # this effectively test the RLock vs Lock
     bear_store = BearReactive(bears)