feat: allows tasks to have their own kernel context

maartenbreddels · maartenbreddels · commit 3a97835ac485 · 2025-12-01T14:18:00.000+01:00
Adding this so we can use solara tasks and reactive variables in
fastapi endpoint, which is only async.
diff --git a/solara/server/kernel_context.py b/solara/server/kernel_context.py
@@ -17,7 +17,7 @@
 import time
 import typing
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 import ipywidgets as widgets
 import reacton
@@ -39,6 +39,11 @@ class Local(threading.local):
 
 
 local = Local()
+# same idea, but for `async with ...`
+if typing.TYPE_CHECKING:
+    async_stack = contextvars.ContextVar[Union[Tuple[Union[None, "VirtualKernelContext"], ...], None]](name="async_stack", default=None)
+else:
+    async_stack = contextvars.ContextVar("async_stack", default=None)
 
 
 class PageStatus(enum.Enum):
@@ -100,6 +105,23 @@ def display(self, *args):
     def on_close(self, f: Callable[[], None]):
         self._on_close_callbacks.append(f)
 
+    async def __aenter__(self):
+        stack = async_stack.get()
+        if stack is None:
+            stack = ()
+        key = get_current_thread_key()
+        async_stack.set(stack + (current_context.get(key, None),))
+        new_key = get_current_thread_key()
+        current_context[new_key] = self
+
+    async def __aexit__(self, *args):
+        key = get_current_thread_key()
+        assert local.kernel_context_stack is not None
+        stack = async_stack.get()
+        assert stack is not None
+        current_context[key] = stack[-1]
+        async_stack.set(stack[:-1])
+
     def __enter__(self):
         if local.kernel_context_stack is None:
             local.kernel_context_stack = []
@@ -364,6 +386,7 @@ def create_dummy_context():
 
 
 def get_current_thread_key() -> str:
+    # consider renaming this to get_current_context_key
     if not solara.server.settings.kernel.threaded:
         if async_context_id is not None:
             try:
@@ -375,6 +398,13 @@ def get_current_thread_key() -> str:
     else:
         thread = threading.current_thread()
         key = get_thread_key(thread)
+        # this signals we are using `async with context`, which means we are interested in task-local context
+        stack = async_stack.get()
+        if stack is not None and len(stack) > 0:
+            current_task = asyncio.current_task()
+            if current_task is not None:
+                task_key = current_task.get_name()
+                key = f"{key}-task:{task_key}"
     return key
 
 
diff --git a/solara/settings.py b/solara/settings.py
@@ -57,6 +57,7 @@ class MainSettings(BaseSettings):
     allow_reactive_boolean: bool = True
     # TODO: also change default_container in solara/components/__init__.py
     default_container: Optional[str] = "Column"
+    allow_global_context: bool = True
 
     class Config:
         env_prefix = "solara_"
diff --git a/solara/tasks.py b/solara/tasks.py
@@ -1,3 +1,4 @@
+import contextvars
 import sys
 import abc
 import asyncio
@@ -267,7 +268,16 @@ def cancel():
 
         if self.run_in_thread:
             thread_event_loop = asyncio.new_event_loop()
-            self.current_task = current_task = thread_event_loop.create_task(self._async_run(call_event_loop, future, args, kwargs))
+
+            def create_task():
+                # remove the stack, since this thread starts with a fresh stack
+                import solara.server.kernel_context
+
+                solara.server.kernel_context.async_stack.set(None)
+                return thread_event_loop.create_task(self._async_run(call_event_loop, future, args, kwargs))
+
+            new_context = contextvars.copy_context()
+            self.current_task = current_task = new_context.run(create_task)
 
             def runs_in_thread():
                 try:
@@ -298,7 +308,7 @@ def runs_in_thread():
                     raise
 
             self._result.value = TaskResult[R](latest=self._last_value, _state=TaskState.STARTING)
-            thread = threading.Thread(target=runs_in_thread, daemon=True)
+            thread = threading.Thread(target=runs_in_thread, daemon=True, name=f"TaskAsyncio-{self.function.__name__}")
             thread.start()
         else:
             self.current_task = current_task = asyncio.create_task(self._async_run(call_event_loop, future, args, kwargs))
@@ -322,7 +332,6 @@ async def _async_run(self, call_event_loop: asyncio.AbstractEventLoop, future: a
 
         task_for_this_call = _get_current_task()
         assert task_for_this_call is not None
-
         if self.is_current():
             self._result.value = TaskResult[R](latest=self._last_value, _state=TaskState.STARTING)
 
diff --git a/solara/toestand.py b/solara/toestand.py
@@ -313,6 +313,10 @@ def clear(self):
 
     def set(self, value: S):
         scope_dict, scope_id = self._get_dict()
+        if not solara.settings.main.allow_global_context and scope_id == "global":
+            raise RuntimeError(
+                f"No kernel context found, and global context is not allowed for task, context key was {solara.server.kernel_context.get_current_thread_key()}"
+            )
         old = self.get()
         if self.equals(old, value):
             return
diff --git a/tests/unit/async_test.py b/tests/unit/async_test.py
@@ -0,0 +1,171 @@
+import asyncio
+from typing import Tuple
+import solara
+from solara.server import kernel_context
+from solara.server import kernel
+from unittest import mock
+import solara.lab
+
+reactive = solara.reactive(0)
+test_async_task_setting = solara.reactive(0)
+tasks = []  # always keep a reference to an asyncio.Task
+
+
+@solara.lab.task  # (prefer_threaded=True)
+async def multiply_by(value: int):
+    result = reactive.value * value
+
+    # we also test that a task created in a new thread
+    async def set_in_task():
+        test_async_task_setting.value = result
+
+    task = asyncio.create_task(set_in_task())
+    tasks.append(task)
+    await asyncio.sleep(0.1)  # give the task a chance to run
+    await task  # wait for the task to finish
+    return result
+
+
+@mock.patch("solara._using_solara_server", return_value=True)
+async def test_async_kernels_basic(_):
+    assert _() is True
+    kernel1 = kernel.Kernel()
+    kernel2 = kernel.Kernel()
+    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")
+
+    values = solara.Reactive[Tuple[int, ...]]((1,))
+
+    async def task1():
+        async with context1:
+            for i in range(99):
+                values.value = values.value + (len(values.value),)
+                await asyncio.sleep(0.01)
+
+    async def task2():
+        async with context2:
+            for i in range(99):
+                values.value = values.value + (len(values.value),)
+                await asyncio.sleep(0.01)
+
+    # await asyncio.gather(asyncio.create_task(task1(), name="test-task1"), asyncio.create_task(task2(), name="test-task2"))
+    await asyncio.gather(task1(), task2())
+    assert values.value == (1,)
+    with context1:
+        assert len(values.value) == 100
+        assert values.value[-1] == 99
+    with context2:
+        assert len(values.value) == 100
+        assert values.value[-1] == 99
+    assert values.value == (1,)
+    async with context1:
+        assert len(values.value) == 100
+        assert values.value[-1] == 99
+    async with context2:
+        assert len(values.value) == 100
+        assert values.value[-1] == 99
+    assert values.value == (1,)
+
+
+@mock.patch("solara._using_solara_server", return_value=True)
+async def test_async_kernels_complex(_):
+    assert _() is True
+    event1 = asyncio.Event()  # after event, global is 1
+    event2 = asyncio.Event()  # after event, global is still 1
+    event3 = asyncio.Event()  # after event, global is 2
+    event4 = asyncio.Event()  # after event, global is 3
+    event5 = asyncio.Event()  # after event, global is 3
+    kernel1 = kernel.Kernel()
+    kernel2 = kernel.Kernel()
+    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")
+
+    main_thread_key = kernel_context.get_current_thread_key()
+
+    async def task1():
+        # global default scope
+        reactive.value = 1
+        event1.set()
+        async with context1:
+            # kernel scope
+            reactive.value = 100
+            assert reactive.value == 100
+            await event3.wait()
+            assert reactive.value == 100
+            multiply_by(3)  # result should be 300
+
+    async def task2():
+        await event2.wait()
+        # global default scope
+        assert reactive.value == 1
+        reactive.value = 2
+        multiply_by(8)  # result should be 16
+        event3.set()
+        assert kernel_context.get_current_thread_key() == main_thread_key
+        async with context2:
+            assert main_thread_key in kernel_context.get_current_thread_key() and len(kernel_context.get_current_thread_key()) > len(main_thread_key)
+            # kernel scope
+            assert reactive.value == 0  # still the default value
+            reactive.value = 200
+            event4.set()
+            multiply_by(3)  # result should be 600
+
+        assert kernel_context.get_current_thread_key() == main_thread_key
+        await event5.wait()
+        assert reactive.value == 3
+
+    async def test():
+        await event1.wait()
+        assert reactive.value == 1
+        event2.set()
+        await event2.wait()
+        # still global default scope
+        assert reactive.value == 1
+        await event3.wait()
+        assert reactive.value == 2
+        await event4.wait()
+        await multiply_by.current_future  # type: ignore
+        reactive.value = 3
+        event5.set()
+
+    await asyncio.gather(task1(), task2(), test())
+
+    with context1:
+        assert reactive.value == 100
+    with context2:
+        assert reactive.value == 200
+    assert reactive.value == 3
+
+    # checking task results
+    while not multiply_by.result.finished:
+        await asyncio.sleep(0.1)
+    assert multiply_by.result.value == 16
+    assert test_async_task_setting.value == 16
+    with context1:
+        while not multiply_by.result.finished:
+            await asyncio.sleep(0.1)
+        assert multiply_by.result.value == 300
+        assert test_async_task_setting.value == 300
+    with context2:
+        while not multiply_by.result.finished:
+            await asyncio.sleep(0.1)
+        assert multiply_by.result.value == 600
+        assert test_async_task_setting.value == 600
+
+
+@mock.patch("solara._using_solara_server", return_value=True)
+async def test_async_kernels_task(_):
+    assert _() is True
+    kernel1 = kernel.Kernel()
+    context1 = kernel_context.VirtualKernelContext(id="toestand-1", kernel=kernel1, session_id="session-1")
+
+    main_thread_key = kernel_context.get_current_thread_key()
+
+    async with context1:
+        assert main_thread_key != kernel_context.get_current_thread_key()
+        assert "task" in kernel_context.get_current_thread_key()
+        reactive.value = 100
+        multiply_by(3)  # result should be 300
+        assert reactive.value == 100
+        await multiply_by.current_future  # type: ignore
+        assert test_async_task_setting.value == 300, "if this fails, the solara task was using the global context"
