diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 8caa714c..4acf8c38 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -74,6 +74,62 @@ def _sanitize_surrogates(text: str) -> str: return text.encode("utf-16-le", errors="surrogatepass").decode("utf-16-le", errors="replace") +def _signal_name(signum: int) -> str: + with suppress(ValueError): + return signal.Signals(signum).name + return f"signal {signum}" + + +def _install_gateway_shutdown_handlers( + loop: asyncio.AbstractEventLoop, + shutdown_event: asyncio.Event, + tasks: list[asyncio.Task], + print_status: Callable[[str], None], +) -> Callable[[], None]: + """Install foreground gateway signal handlers and return a restore callback.""" + loop_signals: list[int] = [] + previous_handlers: list[tuple[int, Any]] = [] + shutdown_requested = False + + def request_shutdown(signum: int) -> None: + nonlocal shutdown_requested + sig_name = _signal_name(signum) + if shutdown_requested: + logger.warning("Forcing gateway shutdown after repeated {}", sig_name) + for task in tasks: + if not task.done(): + task.cancel() + return + shutdown_requested = True + logger.info("Gateway shutdown requested by {}", sig_name) + print_status("\nShutting down... Press Ctrl+C again to force.") + shutdown_event.set() + + for signum in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(signum, request_shutdown, signum) + except (NotImplementedError, RuntimeError, ValueError): + try: + previous = signal.getsignal(signum) + signal.signal(signum, lambda sig, _frame: request_shutdown(sig)) + except (RuntimeError, ValueError): + logger.debug("Could not install gateway handler for {}", _signal_name(signum)) + continue + previous_handlers.append((signum, previous)) + else: + loop_signals.append(signum) + + def restore() -> None: + for signum in loop_signals: + with suppress(NotImplementedError, RuntimeError, ValueError): + loop.remove_signal_handler(signum) + for signum, handler in previous_handlers: + with suppress(RuntimeError, ValueError): + signal.signal(signum, handler) + + return restore + + class SafeFileHistory(FileHistory): """FileHistory subclass that sanitizes surrogate characters on write. @@ -1106,6 +1162,16 @@ def _run_gateway( async def run(): tasks: list[asyncio.Task] = [] + shutdown_task: asyncio.Task | None = None + runtime_tasks: asyncio.Future | None = None + runtime_tasks_drained = False + shutdown_event = asyncio.Event() + restore_shutdown_handlers = _install_gateway_shutdown_handlers( + asyncio.get_running_loop(), + shutdown_event, + tasks, + console.print, + ) try: await cron.start() tasks = [ @@ -1122,7 +1188,20 @@ def _run_gateway( _open_browser_when_ready(), name="nanobot-open-browser", )) - await asyncio.gather(*tasks) + runtime_tasks = asyncio.gather(*tasks) + shutdown_task = asyncio.create_task( + shutdown_event.wait(), + name="nanobot-gateway-shutdown", + ) + done, _pending = await asyncio.wait( + {runtime_tasks, shutdown_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + if runtime_tasks in done: + runtime_tasks_drained = True + await runtime_tasks + elif runtime_tasks is not None: + runtime_tasks.cancel() except KeyboardInterrupt: console.print("\nShutting down...") except Exception: @@ -1131,20 +1210,30 @@ def _run_gateway( console.print("\n[red]Error: Gateway crashed unexpectedly[/red]") console.print(traceback.format_exc()) finally: - cron.stop() - agent.stop() - for task in tasks: - if not task.done(): - task.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - await channels.stop_all() - # Flush all cached sessions to durable storage before exit. - # This prevents data loss on filesystems with write-back - # caching (rclone VFS, NFS, FUSE mounts, etc.). - flushed = agent.sessions.flush_all() - if flushed: - logger.info("Shutdown: flushed {} session(s) to disk", flushed) + try: + if shutdown_task and not shutdown_task.done(): + shutdown_task.cancel() + with suppress(asyncio.CancelledError): + await shutdown_task + cron.stop() + agent.stop() + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + if runtime_tasks is not None and not runtime_tasks_drained: + with suppress(asyncio.CancelledError, Exception): + await runtime_tasks + await channels.stop_all() + # Flush all cached sessions to durable storage before exit. + # This prevents data loss on filesystems with write-back + # caching (rclone VFS, NFS, FUSE mounts, etc.). + flushed = agent.sessions.flush_all() + if flushed: + logger.info("Shutdown: flushed {} session(s) to disk", flushed) + finally: + restore_shutdown_handlers() asyncio.run(run()) diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index 7825bc5d..0e92fdc8 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -2,6 +2,8 @@ import asyncio import json import re import shutil +import signal +from contextlib import suppress from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -9,6 +11,7 @@ import pytest from typer.testing import CliRunner from nanobot.bus.events import InboundMessage, OutboundMessage +from nanobot.cli import commands as cli_commands from nanobot.cli.commands import app from nanobot.config.schema import Config from nanobot.cron.service import CronJobSkippedError @@ -58,6 +61,54 @@ class _StopGatewayError(RuntimeError): pass +def test_gateway_signal_handler_first_signal_stops_and_second_forces() -> None: + class _FakeLoop: + def __init__(self) -> None: + self.handlers: dict[int, tuple[object, tuple[object, ...]]] = {} + self.removed: list[int] = [] + + def add_signal_handler(self, signum, callback, *args) -> None: + self.handlers[int(signum)] = (callback, args) + + def remove_signal_handler(self, signum) -> bool: + self.removed.append(int(signum)) + self.handlers.pop(int(signum), None) + return True + + async def _run() -> None: + loop = _FakeLoop() + shutdown_event = asyncio.Event() + never = asyncio.Event() + task = asyncio.create_task(never.wait()) + output: list[str] = [] + + restore = cli_commands._install_gateway_shutdown_handlers( + loop, shutdown_event, [task], output.append, + ) + try: + callback, args = loop.handlers[int(signal.SIGINT)] + assert callable(callback) + + callback(*args) + assert shutdown_event.is_set() + assert output == ["\nShutting down... Press Ctrl+C again to force."] + assert not task.done() + + callback(*args) + await asyncio.sleep(0) + assert task.cancelled() + finally: + restore() + task.cancel() + with suppress(asyncio.CancelledError): + await task + + assert int(signal.SIGINT) in loop.removed + assert int(signal.SIGTERM) in loop.removed + + asyncio.run(_run()) + + @pytest.fixture def mock_paths(): """Mock config/workspace paths for test isolation.""" @@ -2057,6 +2108,124 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup( assert seen["cron_stopped"] is True +def test_gateway_shutdown_event_exits_forever_runtime_tasks( + monkeypatch, + tmp_path: Path, +) -> None: + config_file = _write_instance_config(tmp_path) + config = Config() + config.gateway.port = 18791 + seen: dict[str, object] = {} + + class _FakeSessionManager: + def flush_all(self) -> int: + return 0 + + class _FakeAgentLoop: + @classmethod + def from_config(cls, config, bus=None, **extra): + return cls(**extra) + + def __init__(self, **_kwargs) -> None: + self.model = "test-model" + self.provider = object() + self.sessions = _FakeSessionManager() + + def llm_runtime(self) -> None: + return None + + async def run(self) -> None: + try: + await asyncio.Event().wait() + finally: + seen["agent_task_cleaned_up"] = True + + async def close_mcp(self) -> None: + raise AssertionError("gateway must not close MCP from the outer task") + + def stop(self) -> None: + seen["agent_stopped"] = True + + class _FakeChannelManager: + def __init__(self, _config, _bus, **_kwargs) -> None: + self.enabled_channels = ["websocket"] + + async def start_all(self) -> None: + try: + await asyncio.Event().wait() + finally: + seen["channel_task_cleaned_up"] = True + + async def stop_all(self) -> None: + seen["channels_stopped"] = True + + class _FakeCronService: + def __init__(self, _store_path: Path) -> None: + self.on_job = None + + async def start(self) -> None: + return None + + def stop(self) -> None: + seen["cron_stopped"] = True + + def status(self) -> dict[str, int]: + return {"jobs": 0} + + def register_system_job(self, _job) -> None: + return None + + class _FakeServer: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + async def serve_forever(self) -> None: + await asyncio.Event().wait() + + async def _fake_start_server(_handler, _host: str, _port: int): + return _FakeServer() + + def _fake_install_shutdown_handlers(_loop, event, _tasks, _print_status): + async def _trigger_shutdown() -> None: + await asyncio.sleep(0) + event.set() + + asyncio.create_task(_trigger_shutdown()) + + def _restore() -> None: + seen["shutdown_handlers_restored"] = True + + return _restore + + _patch_cli_command_runtime( + monkeypatch, + config, + message_bus=lambda: object(), + session_manager=lambda _workspace: object(), + ) + monkeypatch.setattr("nanobot.cli.commands.AgentLoop", _FakeAgentLoop) + monkeypatch.setattr("nanobot.channels.manager.ChannelManager", _FakeChannelManager) + monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCronService) + monkeypatch.setattr("asyncio.start_server", _fake_start_server) + monkeypatch.setattr( + "nanobot.cli.commands._install_gateway_shutdown_handlers", + _fake_install_shutdown_handlers, + ) + + result = runner.invoke(app, ["gateway", "--config", str(config_file)]) + + assert result.exit_code == 0 + assert seen["agent_stopped"] is True + assert seen["agent_task_cleaned_up"] is True + assert seen["channel_task_cleaned_up"] is True + assert seen["channels_stopped"] is True + assert seen["cron_stopped"] is True + assert seen["shutdown_handlers_restored"] is True + + def test_serve_uses_api_config_defaults_and_workspace_override( monkeypatch, tmp_path: Path ) -> None: