Coverage for src/lilbee/runtime/asyncio_loop.py: 100%

1"""Process-lifetime background asyncio loop for TUI workers. 

2 

3One loop on a daemon thread. Every @work(thread=True) worker and every provider 

4backend call (chat, stream, embed, rerank) runs its coroutines here, from the 

5TUI, the CLI and the server alike. 

6""" 

7 

8from __future__ import annotations 

9 

10import asyncio 

11import atexit 

12import concurrent.futures 

13import contextlib 

14import threading 

15from collections.abc import Coroutine 

16from typing import Any, TypeVar 

17 

18T = TypeVar("T") 

19 

20# concurrent.futures raises this exact RuntimeError message when submitting to 

21# a shutdown executor (Python 3.11+). There is no dedicated exception class to 

22# catch, so callers have to string-match the message. The constant + helper 

23# below give every caller one place to depend on instead of duplicating the 

24# string + isinstance check inline. 

25_EXECUTOR_SHUTDOWN_MSG = "cannot schedule new futures after shutdown" 

26# Seconds the pending tasks get to unwind when the loop shuts down. 

27_DRAIN_BUDGET_S = 10.0 

28 

29 

30def is_executor_shutdown(exc: BaseException) -> bool: 

31 """True if ``exc`` is the concurrent.futures shutdown-race ``RuntimeError``.""" 

32 return isinstance(exc, RuntimeError) and _EXECUTOR_SHUTDOWN_MSG in str(exc) 

33 

34 

35_loop: asyncio.AbstractEventLoop | None = None 

36_thread: threading.Thread | None = None 

37_lock = threading.Lock() 

38_atexit_registered = False 

39 

40 

41def get_loop() -> asyncio.AbstractEventLoop: 

42 """Return the background loop, starting it on a daemon thread if needed.""" 

43 global _loop, _thread, _atexit_registered 

44 with _lock: 

45 if _loop is not None and not _loop.is_closed(): 

46 return _loop 

47 loop = asyncio.new_event_loop() 

48 thread = threading.Thread( 

49 target=loop.run_forever, 

50 name="lilbee-bg-loop", 

51 daemon=True, 

52 ) 

53 thread.start() 

54 _loop = loop 

55 _thread = thread 

56 if not _atexit_registered: 

57 # Register once per process; shutdown is idempotent, so restarting 

58 # the loop later doesn't need a second registration. 

59 atexit.register(shutdown) 

60 _atexit_registered = True 

61 return loop 

62 

63 

64def run(coro: Coroutine[Any, Any, T]) -> T: 

65 """Submit *coro* to the background loop from any thread; block for result. 

66 

67 Exceptions raised inside *coro* propagate unchanged, including 

68 asyncio.CancelledError. 

69 """ 

70 loop = get_loop() 

71 try: 

72 return asyncio.run_coroutine_threadsafe(coro, loop).result() 

73 except concurrent.futures.CancelledError as exc: 

74 # run_coroutine_threadsafe re-raises cancellation as the concurrent 

75 # flavour; rewrap so `except asyncio.CancelledError` still matches. 

76 raise asyncio.CancelledError(*exc.args) from None 

77 

78 

79def shutdown(drain_budget_s: float = _DRAIN_BUDGET_S) -> None: 

80 """Cancel pending tasks and wait *drain_budget_s* for them, then stop the loop. Idempotent.""" 

81 global _loop, _thread 

82 with _lock: 

83 loop, _loop = _loop, None 

84 thread, _thread = _thread, None 

85 if loop is None or loop.is_closed(): 

86 return 

87 # Best-effort drain; always stop the loop even if drain raised. 

88 with contextlib.suppress(Exception): 

89 asyncio.run_coroutine_threadsafe(_drain(loop), loop).result(timeout=drain_budget_s) 

90 loop.call_soon_threadsafe(loop.stop) 

91 if thread is not None: 

92 thread.join(timeout=10.0) 

93 loop.close() 

94 

95 

96async def _drain(loop: asyncio.AbstractEventLoop) -> None: 

97 pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task()] 

98 for task in pending: 

99 task.cancel() 

100 if pending: 

101 await asyncio.gather(*pending, return_exceptions=True) 

102 # Give scheduled close callbacks a chance to run before we stop the loop. 

103 await asyncio.sleep(0.05)