fix(grpc): support signal shutdown on Windows

python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_signal_utils.py

@@ -0,0 +1,45 @@
+import asyncio
+import logging
+import signal
+from collections.abc import Sequence
+from types import FrameType
+from typing import Any
+
+logger = logging.getLogger("autogen_core")
+
+
+async def wait_for_signal(signals: Sequence[signal.Signals]) -> None:
+    """Wait for one of *signals* on event loops with or without signal support."""
+    loop = asyncio.get_running_loop()
+    shutdown_event = asyncio.Event()
+    unique_signals = tuple(dict.fromkeys(signals))
+    loop_signals: list[signal.Signals] = []
+    previous_handlers: dict[signal.Signals, Any] = {}
+
+    def signal_handler() -> None:
+        logger.info("Received exit signal, shutting down gracefully...")
+        shutdown_event.set()
+
+    def fallback_signal_handler(_signum: int, _frame: FrameType | None) -> None:
+        loop.call_soon_threadsafe(signal_handler)
+
+    try:
+        try:
+            for sig in unique_signals:
+                loop.add_signal_handler(sig, signal_handler)
+                loop_signals.append(sig)
+        except NotImplementedError:
+            # Windows' default event loop does not implement add_signal_handler.
+            for sig in loop_signals:
+                loop.remove_signal_handler(sig)
+            loop_signals.clear()
+
+            for sig in unique_signals:
+                previous_handlers[sig] = signal.signal(sig, fallback_signal_handler)
+
+        await shutdown_event.wait()
+    finally:
+        for sig in loop_signals:
+            loop.remove_signal_handler(sig)
+        for sig, previous_handler in previous_handlers.items():
+            signal.signal(sig, previous_handler)

python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py

@@ -60,6 +60,7 @@
 
 from . import _constants
 from ._constants import GRPC_IMPORT_ERROR_STR
+from ._signal_utils import wait_for_signal
 from ._type_helpers import ChannelArgumentType
 from .protos import agent_worker_pb2, agent_worker_pb2_grpc, cloudevent_pb2
 
@@ -331,20 +332,7 @@ async def stop(self) -> None:
 
     async def stop_when_signal(self, signals: Sequence[signal.Signals] = (signal.SIGTERM, signal.SIGINT)) -> None:
         """Stop the runtime when a signal is received."""
-        loop = asyncio.get_running_loop()
-        shutdown_event = asyncio.Event()
-
-        def signal_handler() -> None:
-            logger.info("Received exit signal, shutting down gracefully...")
-            shutdown_event.set()
-
-        for sig in signals:
-            loop.add_signal_handler(sig, signal_handler)
-
-        # Wait for the signal to trigger the shutdown event.
-        await shutdown_event.wait()
-
-        # Stop the runtime.
+        await wait_for_signal(signals)
         await self.stop()
 
     @property

python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host.py

@@ -4,6 +4,7 @@
 from typing import Optional, Sequence
 
 from ._constants import GRPC_IMPORT_ERROR_STR
+from ._signal_utils import wait_for_signal
 from ._type_helpers import ChannelArgumentType
 from ._worker_runtime_host_servicer import GrpcWorkerAgentRuntimeHostServicer
 
@@ -56,18 +57,7 @@ async def stop_when_signal(
         if self._serve_task is None:
             raise RuntimeError("Host runtime is not started.")
         # Set up signal handling for graceful shutdown.
-        loop = asyncio.get_running_loop()
-        shutdown_event = asyncio.Event()
-
-        def signal_handler() -> None:
-            logger.info("Received exit signal, shutting down gracefully...")
-            shutdown_event.set()
-
-        for sig in signals:
-            loop.add_signal_handler(sig, signal_handler)
-
-        # Wait for the signal to trigger the shutdown event.
-        await shutdown_event.wait()
+        await wait_for_signal(signals)
 
         # Shutdown the server.
         await self.stop(grace=grace)

python/packages/autogen-ext/tests/test_grpc_signal_utils.py

@@ -0,0 +1,66 @@
+import asyncio
+import signal
+from collections.abc import Callable
+from types import FrameType
+from typing import Any
+
+import pytest
+from autogen_ext.runtimes.grpc import GrpcWorkerAgentRuntime
+from autogen_ext.runtimes.grpc._signal_utils import wait_for_signal
+
+
+@pytest.mark.asyncio
+async def test_wait_for_signal_uses_and_removes_loop_handlers(monkeypatch: pytest.MonkeyPatch) -> None:
+    loop = asyncio.get_running_loop()
+    installed_handlers: dict[signal.Signals, Callable[[], None]] = {}
+    removed_signals: list[signal.Signals] = []
+
+    def add_signal_handler(sig: signal.Signals, callback: Callable[[], None], *args: Any) -> None:
+        installed_handlers[sig] = callback
+
+    def remove_signal_handler(sig: signal.Signals) -> bool:
+        removed_signals.append(sig)
+        return True
+
+    monkeypatch.setattr(loop, "add_signal_handler", add_signal_handler)
+    monkeypatch.setattr(loop, "remove_signal_handler", remove_signal_handler)
+
+    wait_task = asyncio.create_task(wait_for_signal((signal.SIGINT,)))
+    await asyncio.sleep(0)
+    installed_handlers[signal.SIGINT]()
+    await wait_task
+
+    assert removed_signals == [signal.SIGINT]
+
+
+@pytest.mark.asyncio
+async def test_worker_runtime_falls_back_and_restores_handlers(monkeypatch: pytest.MonkeyPatch) -> None:
+    loop = asyncio.get_running_loop()
+    installed_handlers: dict[signal.Signals, signal.Handlers | Callable[[int, FrameType | None], Any]] = {}
+    previous_handler = signal.SIG_DFL
+
+    def unsupported_signal_handler(sig: signal.Signals, callback: Callable[[], None], *args: Any) -> None:
+        raise NotImplementedError
+
+    def replace_signal_handler(
+        sig: signal.Signals, handler: signal.Handlers | Callable[[int, FrameType | None], Any]
+    ) -> signal.Handlers | Callable[[int, FrameType | None], Any]:
+        old_handler = installed_handlers.get(sig, previous_handler)
+        installed_handlers[sig] = handler
+        return old_handler
+
+    monkeypatch.setattr(loop, "add_signal_handler", unsupported_signal_handler)
+    monkeypatch.setattr(signal, "signal", replace_signal_handler)
+
+    runtime = GrpcWorkerAgentRuntime(host_address="unused")
+    runtime._running = True  # type: ignore[reportPrivateUsage]
+    wait_task = asyncio.create_task(runtime.stop_when_signal((signal.SIGINT, signal.SIGINT)))
+    await asyncio.sleep(0)
+
+    fallback_handler = installed_handlers[signal.SIGINT]
+    assert callable(fallback_handler)
+    fallback_handler(signal.SIGINT, None)
+    await wait_task
+
+    assert not runtime._running  # type: ignore[reportPrivateUsage]
+    assert installed_handlers[signal.SIGINT] is previous_handler