Add .on_end() callback registration for before the WS is closed

`reactive.isolate()` before calling callbacks

Author: schloerke

shiny/express/_stub_session.py

@@ -58,6 +58,12 @@ async def close(self, code: int = 1001) -> None:
     def _is_hidden(self, name: str) -> bool:
         return False
 
+    def on_end(
+        self,
+        fn: Callable[[], None] | Callable[[], Awaitable[None]],
+    ) -> Callable[[], None]:
+        return lambda: None
+
     def on_ended(
         self,
         fn: Callable[[], None] | Callable[[], Awaitable[None]],

shiny/session/_session.py

@@ -211,6 +211,25 @@ async def close(self, code: int = 1001) -> None:
     @abstractmethod
     def _is_hidden(self, name: str) -> bool: ...
 
+    @abstractmethod
+    def on_end(
+        self,
+        fn: Callable[[], None] | Callable[[], Awaitable[None]],
+    ) -> Callable[[], None]:
+        """
+        Registers a function to be called hopefully before the client is disconnected.
+
+        Parameters
+        ----------
+        fn
+            The function to call.
+
+        Returns
+        -------
+        :
+            A function that can be used to cancel the registration.
+        """
+
     @add_example()
     @abstractmethod
     def on_ended(
@@ -559,7 +578,9 @@ def __init__(
         self._outbound_message_queues = OutBoundMessageQueues()
 
         self._file_upload_manager: FileUploadManager = FileUploadManager()
+        self._on_end_callbacks = _utils.AsyncCallbacks()
         self._on_ended_callbacks = _utils.AsyncCallbacks()
+        self._has_run_session_end_tasks: bool = False
         self._has_run_session_ended_tasks: bool = False
         self._downloads: dict[str, DownloadInfo] = {}
         self._dynamic_routes: dict[str, DynamicRouteHandler] = {}
@@ -576,20 +597,37 @@ def _register_session_ended_callbacks(self) -> None:
         # Clear file upload directories, if present
         self.on_ended(self._file_upload_manager.rm_upload_dir)
 
+    async def _run_session_end_tasks(self) -> None:
+        if self._has_run_session_end_tasks:
+            return
+        self._has_run_session_end_tasks = True
+
+        try:
+            with isolate():
+                await self._on_end_callbacks.invoke()
+        except Exception as e:
+            warnings.warn(
+                f"Error running session end tasks: {str(e)}",
+                SessionWarning,
+                stacklevel=2,
+            )
+
     async def _run_session_ended_tasks(self) -> None:
         if self._has_run_session_ended_tasks:
             return
         self._has_run_session_ended_tasks = True
 
         try:
-            await self._on_ended_callbacks.invoke()
+            with isolate():
+                await self._on_ended_callbacks.invoke()
         finally:
             self.app._remove_session(self)
 
     def is_stub_session(self) -> Literal[False]:
         return False
 
     async def close(self, code: int = 1001) -> None:
+        await self._run_session_end_tasks()
         await self._conn.close(code, None)
         await self._run_session_ended_tasks()
 
@@ -713,6 +751,7 @@ def verify_state(expected_state: ConnectionState) -> None:
                 finally:
                     await self.close()
             finally:
+                await self._run_session_end_tasks()
                 await self._run_session_ended_tasks()
 
     def _manage_inputs(self, data: dict[str, object]) -> None:
@@ -1079,8 +1118,15 @@ def _decrement_busy_count(self) -> None:
             self._send_message_sync({"busy": "idle"})
 
     # ==========================================================================
-    # On session ended
+    # On session end / ended
     # ==========================================================================
+
+    def on_end(
+        self,
+        fn: Callable[[], None] | Callable[[], Awaitable[None]],
+    ) -> Callable[[], None]:
+        return self._on_end_callbacks.register(wrap_async(fn))
+
     def on_ended(
         self,
         fn: Callable[[], None] | Callable[[], Awaitable[None]],
@@ -1221,6 +1267,12 @@ def __init__(self, root_session: Session, ns: ResolvedId) -> None:
     def _is_hidden(self, name: str) -> bool:
         return self._root_session._is_hidden(name)
 
+    def on_end(
+        self,
+        fn: Callable[[], None] | Callable[[], Awaitable[None]],
+    ) -> Callable[[], None]:
+        return self._root_session.on_end(fn)
+
     def on_ended(
         self,
         fn: Callable[[], None] | Callable[[], Awaitable[None]],