feat: on_kernel_start triggers callback on virtual kernel start

solara/lab/__init__.py

@@ -1,5 +1,6 @@
 # isort: skip_file
 from .components import *  # noqa: F401, F403
+from ..server.kernel_context import on_kernel_start  # noqa: F401
 
 
 def __getattr__(name):

solara/server/app.py

@@ -375,6 +375,7 @@ def on_msg(msg):
             context = kernel_context.get_current_context()
             path = data.get("path", "")
             with context:
+                context.restart()
                 load_app_widget(context.state, app, path)
                 comm.send({"method": "finished"})
         else:

solara/server/kernel_context.py

@@ -36,6 +36,13 @@ class PageStatus(enum.Enum):
     CLOSED = "closed"
 
 
+_on_kernel_start_callbacks: List[Callable[[], Optional[Callable[[], None]]]] = []
+
+
+def on_kernel_start(f: Callable[[], Optional[Callable[[], None]]]):
+    _on_kernel_start_callbacks.append(f)
+
+
 @dataclasses.dataclass
 class VirtualKernelContext:
     id: str
@@ -63,10 +70,29 @@ class VirtualKernelContext:
     # only used for testing
     _last_kernel_cull_task: "Optional[asyncio.Future[None]]" = None
     closed_event: threading.Event = dataclasses.field(default_factory=threading.Event)
+    _on_close_callbacks: List[Callable[[], None]] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        with self:
+            for f in _on_kernel_start_callbacks:
+                cleanup = f()
+                if cleanup:
+                    self.on_close(cleanup)
+
+    def restart(self):
+        # should we do this, or maybe close the context and create a new one?
+        with self:
+            for f in reversed(self._on_close_callbacks):
+                f()
+            self._on_close_callbacks.clear()
+            self.__post_init__()
 
     def display(self, *args):
         print(args)  # noqa
 
+    def on_close(self, f: Callable[[], None]):
+        self._on_close_callbacks.append(f)
+
     def __enter__(self):
         if local.kernel_context_stack is None:
             local.kernel_context_stack = []
@@ -81,6 +107,9 @@ def __exit__(self, *args):
 
     def close(self):
         logger.info("Shut down virtual kernel: %s", self.id)
+        with self:
+            for f in reversed(self._on_close_callbacks):
+                f()
         with self:
             if self.app_object is not None:
                 if isinstance(self.app_object, reacton.core._RenderContext):
@@ -284,6 +313,7 @@ def initialize_virtual_kernel(session_id: str, kernel_id: str, websocket: websoc
         kernel = Kernel()
         logger.info("new virtual kernel: %s", kernel_id)
         context = contexts[kernel_id] = VirtualKernelContext(id=kernel_id, session_id=session_id, kernel=kernel, control_sockets=[], widgets={}, templates={})
+
         with context:
             widgets.register_comm_target(kernel)
             appmodule.register_solara_comm_target(kernel)

solara/website/pages/api/__init__.py

@@ -120,6 +120,7 @@
             "confirmation_dialog",
             "menu",
             "input_date",
+            "on_kernel_start",
             "tab",
             "tabs",
         ],

solara/website/pages/api/on_kernel_start.py

@@ -0,0 +1,21 @@
+"""
+# on_kernel_start
+
+Run a function when a virtual kernel (re)starts and optionally run a cleanup function on shutdown.
+
+```python
+def on_kernel_start(f: Callable[[], Optional[Callable[[], None]]]):
+```
+
+`f` will be called on each virtual kernel (re)start. This (usually) happens each time a browser tab connects to the server
+[see solara server for more details](https://solara.dev/docs/understanding/solara-server).
+The (optional) function returned by `f` will be called on kernel shutdown.
+
+Note that the cleanup functions are called in reverse order with respect to the order in which they were registered
+(e.g. the cleanup function of the last call to `on_kernel_start` will be called first on kernel shutdown)
+"""
+
+from . import NoPage
+
+title = "on_kernel_start"
+Page = NoPage