From ba7c07ccf2e81178c107367b048761ab5f4ff4f1 Mon Sep 17 00:00:00 2001 From: imfondof Date: Thu, 2 Apr 2026 16:42:47 +0800 Subject: [PATCH] fix(restart): send completion notice after channel is ready and unify runtime keys --- nanobot/cli/commands.py | 74 ++++++++++++++++++++++++++-- nanobot/command/builtin.py | 3 ++ nanobot/config/runtime_keys.py | 4 ++ tests/cli/test_restart_command.py | 81 ++++++++++++++++++++++++++++++- 4 files changed, 156 insertions(+), 6 deletions(-) create mode 100644 nanobot/config/runtime_keys.py diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index d611c2772..b1e4f056a 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -206,6 +206,57 @@ def _is_exit_command(command: str) -> bool: return command.lower() in EXIT_COMMANDS +def _parse_cli_session(session_id: str) -> tuple[str, str]: + """Split session id into (channel, chat_id).""" + if ":" in session_id: + return session_id.split(":", 1) + return "cli", session_id + + +def _should_show_cli_restart_notice( + restart_notify_channel: str, + restart_notify_chat_id: str, + session_id: str, +) -> bool: + """Return True when CLI should display restart-complete notice.""" + _, cli_chat_id = _parse_cli_session(session_id) + return restart_notify_channel == "cli" and ( + not restart_notify_chat_id or restart_notify_chat_id == cli_chat_id + ) + + +async def _notify_restart_done_when_channel_ready( + *, + bus, + channels, + channel: str, + chat_id: str, + timeout_s: float = 30.0, + poll_s: float = 0.25, +) -> bool: + """Wait for target channel readiness, then publish restart completion.""" + from nanobot.bus.events import OutboundMessage + + if not channel or not chat_id: + return False + if channel not in channels.enabled_channels: + return False + + waited = 0.0 + while waited <= timeout_s: + target = channels.get_channel(channel) + if target and target.is_running: + await bus.publish_outbound(OutboundMessage( + channel=channel, + chat_id=chat_id, + content="Restart completed.", + )) + return True + await asyncio.sleep(poll_s) + waited += poll_s + return False + + async def _read_interactive_input_async() -> str: """Read user input using prompt_toolkit (handles paste, history, display). @@ -598,6 +649,7 @@ def gateway( from nanobot.agent.loop import AgentLoop from nanobot.bus.queue import MessageBus from nanobot.channels.manager import ChannelManager + from nanobot.config.runtime_keys import RESTART_NOTIFY_CHANNEL_ENV, RESTART_NOTIFY_CHAT_ID_ENV from nanobot.cron.service import CronService from nanobot.cron.types import CronJob from nanobot.heartbeat.service import HeartbeatService @@ -696,6 +748,8 @@ def gateway( # Create channel manager channels = ChannelManager(config, bus) + restart_notify_channel = os.environ.pop(RESTART_NOTIFY_CHANNEL_ENV, "").strip() + restart_notify_chat_id = os.environ.pop(RESTART_NOTIFY_CHAT_ID_ENV, "").strip() def _pick_heartbeat_target() -> tuple[str, str]: """Pick a routable channel/chat target for heartbeat-triggered messages.""" @@ -772,6 +826,13 @@ def gateway( try: await cron.start() await heartbeat.start() + if restart_notify_channel and restart_notify_chat_id: + asyncio.create_task(_notify_restart_done_when_channel_ready( + bus=bus, + channels=channels, + channel=restart_notify_channel, + chat_id=restart_notify_chat_id, + )) await asyncio.gather( agent.run(), channels.start_all(), @@ -813,6 +874,7 @@ def agent( from nanobot.agent.loop import AgentLoop from nanobot.bus.queue import MessageBus + from nanobot.config.runtime_keys import RESTART_NOTIFY_CHANNEL_ENV, RESTART_NOTIFY_CHAT_ID_ENV from nanobot.cron.service import CronService config = _load_runtime_config(config, workspace) @@ -853,6 +915,13 @@ def agent( channels_config=config.channels, timezone=config.agents.defaults.timezone, ) + restart_notify_channel = os.environ.pop(RESTART_NOTIFY_CHANNEL_ENV, "").strip() + restart_notify_chat_id = os.environ.pop(RESTART_NOTIFY_CHAT_ID_ENV, "").strip() + + cli_channel, cli_chat_id = _parse_cli_session(session_id) + + if _should_show_cli_restart_notice(restart_notify_channel, restart_notify_chat_id, session_id): + _print_agent_response("Restart completed.", render_markdown=False) # Shared reference for progress callbacks _thinking: ThinkingSpinner | None = None @@ -891,11 +960,6 @@ def agent( _init_prompt_session() console.print(f"{__logo__} Interactive mode (type [bold]exit[/bold] or [bold]Ctrl+C[/bold] to quit)\n") - if ":" in session_id: - cli_channel, cli_chat_id = session_id.split(":", 1) - else: - cli_channel, cli_chat_id = "cli", session_id - def _handle_signal(signum, frame): sig_name = signal.Signals(signum).name _restore_terminal() diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 05d4fc163..f63a1e357 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -9,6 +9,7 @@ import sys from nanobot import __version__ from nanobot.bus.events import OutboundMessage from nanobot.command.router import CommandContext, CommandRouter +from nanobot.config.runtime_keys import RESTART_NOTIFY_CHANNEL_ENV, RESTART_NOTIFY_CHAT_ID_ENV from nanobot.utils.helpers import build_status_content @@ -35,6 +36,8 @@ async def cmd_stop(ctx: CommandContext) -> OutboundMessage: async def cmd_restart(ctx: CommandContext) -> OutboundMessage: """Restart the process in-place via os.execv.""" msg = ctx.msg + os.environ[RESTART_NOTIFY_CHANNEL_ENV] = msg.channel + os.environ[RESTART_NOTIFY_CHAT_ID_ENV] = msg.chat_id async def _do_restart(): await asyncio.sleep(1) diff --git a/nanobot/config/runtime_keys.py b/nanobot/config/runtime_keys.py new file mode 100644 index 000000000..2dc6c9234 --- /dev/null +++ b/nanobot/config/runtime_keys.py @@ -0,0 +1,4 @@ +"""Runtime environment variable keys shared across components.""" + +RESTART_NOTIFY_CHANNEL_ENV = "NANOBOT_RESTART_NOTIFY_CHANNEL" +RESTART_NOTIFY_CHAT_ID_ENV = "NANOBOT_RESTART_NOTIFY_CHAT_ID" diff --git a/tests/cli/test_restart_command.py b/tests/cli/test_restart_command.py index 6efcdad0d..16b3aaa48 100644 --- a/tests/cli/test_restart_command.py +++ b/tests/cli/test_restart_command.py @@ -3,7 +3,9 @@ from __future__ import annotations import asyncio +import os import time +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -35,15 +37,19 @@ class TestRestartCommand: @pytest.mark.asyncio async def test_restart_sends_message_and_calls_execv(self): from nanobot.command.builtin import cmd_restart + from nanobot.config.runtime_keys import RESTART_NOTIFY_CHANNEL_ENV, RESTART_NOTIFY_CHAT_ID_ENV from nanobot.command.router import CommandContext loop, bus = _make_loop() msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart") ctx = CommandContext(msg=msg, session=None, key=msg.session_key, raw="/restart", loop=loop) - with patch("nanobot.command.builtin.os.execv") as mock_execv: + with patch.dict(os.environ, {}, clear=False), \ + patch("nanobot.command.builtin.os.execv") as mock_execv: out = await cmd_restart(ctx) assert "Restarting" in out.content + assert os.environ.get(RESTART_NOTIFY_CHANNEL_ENV) == "cli" + assert os.environ.get(RESTART_NOTIFY_CHAT_ID_ENV) == "direct" await asyncio.sleep(1.5) mock_execv.assert_called_once() @@ -190,3 +196,76 @@ class TestRestartCommand: assert response is not None assert response.metadata == {"render_as": "text"} + + +@pytest.mark.asyncio +async def test_notify_restart_done_waits_until_channel_running() -> None: + from nanobot.bus.queue import MessageBus + from nanobot.cli.commands import _notify_restart_done_when_channel_ready + + bus = MessageBus() + channel = SimpleNamespace(is_running=False) + + class DummyChannels: + enabled_channels = ["feishu"] + + @staticmethod + def get_channel(name: str): + return channel if name == "feishu" else None + + async def _mark_running() -> None: + await asyncio.sleep(0.02) + channel.is_running = True + + marker = asyncio.create_task(_mark_running()) + sent = await _notify_restart_done_when_channel_ready( + bus=bus, + channels=DummyChannels(), + channel="feishu", + chat_id="oc_123", + timeout_s=0.2, + poll_s=0.01, + ) + await marker + + assert sent is True + out = await asyncio.wait_for(bus.consume_outbound(), timeout=0.1) + assert out.channel == "feishu" + assert out.chat_id == "oc_123" + assert out.content == "Restart completed." + + +@pytest.mark.asyncio +async def test_notify_restart_done_times_out_when_channel_not_running() -> None: + from nanobot.bus.queue import MessageBus + from nanobot.cli.commands import _notify_restart_done_when_channel_ready + + bus = MessageBus() + channel = SimpleNamespace(is_running=False) + + class DummyChannels: + enabled_channels = ["feishu"] + + @staticmethod + def get_channel(name: str): + return channel if name == "feishu" else None + + sent = await _notify_restart_done_when_channel_ready( + bus=bus, + channels=DummyChannels(), + channel="feishu", + chat_id="oc_123", + timeout_s=0.05, + poll_s=0.01, + ) + assert sent is False + assert bus.outbound_size == 0 + + +def test_should_show_cli_restart_notice() -> None: + from nanobot.cli.commands import _should_show_cli_restart_notice + + assert _should_show_cli_restart_notice("cli", "direct", "cli:direct") is True + assert _should_show_cli_restart_notice("cli", "", "cli:direct") is True + assert _should_show_cli_restart_notice("cli", "other", "cli:direct") is False + assert _should_show_cli_restart_notice("feishu", "oc_123", "cli:direct") is False