feat: auto hold session as needed

Author: wang0618

pywebio/session/base.py

@@ -29,6 +29,7 @@ class Session:
         next_client_event
         on_task_exception
         register_callback
+        need_keep_alive
 
         defer_call
 
@@ -159,6 +160,9 @@ def defer_call(self, func):
         """设置会话结束时调用的函数。可以用于资源清理。"""
         self.deferred_functions.append(func)
 
+    def need_keep_alive(self) -> bool:
+        raise NotImplementedError
+
 
 def get_session_info_from_headers(headers):
     """从Http请求头中获取会话信息

pywebio/session/coroutinebased.py

@@ -87,14 +87,22 @@ def __init__(self, target, session_info, on_task_command=None, on_session_close=
 
         self._closed = False
 
+        self._need_keep_alive = False
+
         # 当前会话未结束运行(已创建和正在运行的)的协程数量。当 _alive_coro_cnt 变为 0 时，会话结束。
         self._alive_coro_cnt = 1
 
-        main_task = Task(target(), session=self, on_coro_stop=self._on_task_finish)
+        main_task = Task(self._start_main_task(target), session=self, on_coro_stop=self._on_task_finish)
         self.coros[main_task.coro_id] = main_task
 
         self._step_task(main_task)
 
+    async def _start_main_task(self, target):
+        await target()
+        if self.need_keep_alive():
+            from ..session import hold
+            await hold()
+
     def _step_task(self, task, result=None):
         asyncio.get_event_loop().call_soon_threadsafe(partial(task.step, result))
 
@@ -203,6 +211,8 @@ async def callback_coro():
         callback_task.coro.send(None)
         cls.get_current_session().coros[callback_task.coro_id] = callback_task
 
+        self._need_keep_alive = True
+
         return callback_task.coro_id
 
     def run_async(self, coro_obj):
@@ -227,6 +237,9 @@ async def run_asyncio_coroutine(self, coro_obj):
         res = await WebIOFuture(coro=coro_obj)
         return res
 
+    def need_keep_alive(self) -> bool:
+        return self._need_keep_alive
+
 
 class TaskHandler:
     """The handler of coroutine task

pywebio/session/threadbased.py

@@ -91,10 +91,15 @@ def main_task(target):
                 for t in self.threads:
                     if t.is_alive() and t is not threading.current_thread():
                         t.join()
-                try:
-                    self.send_task_command(dict(command='close_session'))
-                except SessionClosedException:
-                    pass
+
+                if self.need_keep_alive():
+                    from ..session import hold
+                    hold()
+                else:
+                    try:
+                        self.send_task_command(dict(command='close_session'))
+                    except SessionClosedException:
+                        pass
                 self._trigger_close_event()
                 self.close()
 
@@ -281,6 +286,10 @@ def register_thread(self, t: threading.Thread):
         event_mq = queue.Queue(maxsize=self.event_mq_maxsize)  # 线程内的用户事件队列
         self.task_mqs[self._get_task_id(t)] = event_mq
 
+    def need_keep_alive(self) -> bool:
+        # if callback thread is activated, then the session need to keep alive
+        return self.callback_thread is not None
+
 
 class ScriptModeSession(ThreadBasedSession):
     """Script mode的会话实现"""

test/1.basic.py

@@ -32,6 +32,8 @@ def test(server_proc: subprocess.Popen, browser: Chrome):
     time.sleep(1)
     template.save_output(browser, '1.basic.html')
 
+    browser.get('http://localhost:8080/')  # to close current session
+
     template.test_defer_call()