fix(cli): don't raise from the worker signal handler so SIGTERM always drains

livekit-agents/livekit/agents/cli/cli.py

@@ -294,25 +294,34 @@ def _run_worker(server: AgentServer, args: proto.CliArgs) -> None:
     devmode = args.dev
     colored_logs = devmode or args.log_format == "colored"
 
-    exit_raised = False
-
-    def _handle_exit(sig: int, frame: FrameType | None) -> None:
-        nonlocal exit_raised
-        if exit_raised:
-            os._exit(1)
-        exit_raised = True
-        raise _ExitCli()
-
-    for sig in HANDLED_SIGNALS:
-        signal.signal(sig, _handle_exit)
-
     setup_logging(args.log_level, devmode=colored_logs, console=False, compact=args.simulation)
 
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
 
     loop.slow_callback_duration = 0.1  # 100ms
 
+    exit_fut: asyncio.Future[None] = loop.create_future()
+    exit_triggered = False
+
+    def _signal_exit() -> None:
+        if not exit_fut.done():
+            exit_fut.set_result(None)
+
+    def _handle_exit(sig: int, frame: FrameType | None) -> None:
+        nonlocal exit_triggered
+        if exit_triggered:
+            os._exit(1)
+        exit_triggered = True
+        # raising from the handler would surface the exception inside whatever
+        # frame the main thread is executing (e.g. a blocking request_fnc inside
+        # a job-request task) instead of at the run_until_complete boundary,
+        # losing the graceful shutdown; schedule the exit on the loop instead
+        loop.call_soon_threadsafe(_signal_exit)
+
+    for sig in HANDLED_SIGNALS:
+        signal.signal(sig, _handle_exit)
+
     async def _worker_run(worker: AgentServer) -> None:
         try:
             await server.run(devmode=devmode, unregistered=False)
@@ -329,7 +338,11 @@ async def _worker_run(worker: AgentServer) -> None:
     try:
         main_task = loop.create_task(_worker_run(server), name="worker_main_task_cli")
         try:
-            loop.run_until_complete(main_task)
+            # exit_fut interrupts the wait on the first signal while main_task
+            # (server.run) keeps running, so the drain below still has a live worker
+            loop.run_until_complete(
+                asyncio.wait([main_task, exit_fut], return_when=asyncio.FIRST_COMPLETED)
+            )
         except _ExitCli:
             pass
 

tests/test_cli_sigterm.py

@@ -0,0 +1,97 @@
+"""Graceful-shutdown signal handling tests for the worker CLI.
+
+A signal handler that raises would surface the exception inside whatever frame
+the main thread happens to be executing — e.g. a user request_fnc doing blocking
+I/O inside a job-request task — instead of at the run_until_complete boundary,
+losing the shutdown entirely (https://github.com/livekit/agents/issues/6724).
+"""
+
+from __future__ import annotations
+
+import signal
+import subprocess
+import sys
+import time
+
+import pytest
+
+pytestmark = pytest.mark.unit
+
+_WORKER_SCRIPT = """
+import asyncio
+import time
+
+from livekit.agents import AgentServer, JobContext
+from livekit.agents.cli import proto
+from livekit.agents.cli.cli import _run_worker
+
+server = AgentServer(
+    ws_url="ws://127.0.0.1:1",  # unreachable: the worker retries forever
+    api_key="devkey",
+    api_secret="devsecret",
+    max_retry=100000,
+    num_idle_processes=0,
+)
+
+
+@server.rtc_session()
+async def entry(ctx: JobContext) -> None:
+    pass
+
+
+_orig_run = server.run
+
+
+async def _run(*args, **kwargs):
+    async def _block_loop() -> None:
+        await asyncio.sleep(0.1)
+        print("BLOCKING", flush=True)
+        # a synchronous call inside a task blocks the event loop on the main
+        # thread, the state a blocking user request_fnc puts the worker in
+        time.sleep(3)
+        print("UNBLOCKED", flush=True)
+
+    asyncio.ensure_future(_block_loop())
+    return await _orig_run(*args, **kwargs)
+
+
+server.run = _run
+
+_run_worker(server, proto.CliArgs(log_level="DEBUG", simulation=True))
+print("CLEAN_EXIT", flush=True)
+"""
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGTERM semantics differ on Windows")
+def test_sigterm_during_blocked_event_loop_shuts_down_worker() -> None:
+    proc = subprocess.Popen(
+        [sys.executable, "-c", _WORKER_SCRIPT],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
+
+    try:
+        assert proc.stdout is not None
+        deadline = time.monotonic() + 30
+        for line in proc.stdout:
+            if "BLOCKING" in line:
+                break
+            if time.monotonic() > deadline:
+                pytest.fail("worker subprocess never reached the blocking section")
+        else:
+            pytest.fail(f"worker subprocess exited early (rc={proc.wait()})")
+
+        # deliver a single SIGTERM while the event loop is blocked; the worker
+        # must still drain and exit once the blocking call returns
+        proc.send_signal(signal.SIGTERM)
+        out, _ = proc.communicate(timeout=30)
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+            proc.communicate()
+
+    assert proc.returncode == 0, f"worker exited with {proc.returncode}:\n{out}"
+    assert "CLEAN_EXIT" in out
+    assert "Task exception was never retrieved" not in out
+    assert "_ExitCli" not in out