mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-04-04 10:22:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(restart): show restart completion with elapsed time across channels

Browse Source
This commit is contained in:
imfondof 2026-04-03 00:44:17 +08:00 committed by Xubin Ren
parent ba7c07ccf2
commit 896d578677
8 changed files with 179 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
nanobot/channels/manager.py
View File

@ -11,6 +11,7 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Config
from nanobot.utils.restart import consume_restart_notice_from_env, format_restart_completed_message
# Retry delays for message sending (exponential backoff: 1s, 2s, 4s)
_SEND_RETRY_DELAYS = (1, 2, 4)
@ -91,9 +92,28 @@ class ChannelManager:
logger.info("Starting {} channel...", name)
tasks.append(asyncio.create_task(self._start_channel(name, channel)))
self._notify_restart_done_if_needed()
# Wait for all to complete (they should run forever)
await asyncio.gather(*tasks, return_exceptions=True)
def _notify_restart_done_if_needed(self) -> None:
"""Send restart completion message when runtime env markers are present."""
notice = consume_restart_notice_from_env()
if not notice:
return
target = self.channels.get(notice.channel)
if not target:
return
asyncio.create_task(self._send_with_retry(
target,
OutboundMessage(
channel=notice.channel,
chat_id=notice.chat_id,
content=format_restart_completed_message(notice.started_at_raw),
),
))
async def stop_all(self) -> None:
"""Stop all channels and the dispatcher."""
logger.info("Stopping all channels...")

85
nanobot/cli/commands.py
View File

@ -37,6 +37,11 @@ from nanobot.cli.stream import StreamRenderer, ThinkingSpinner
from nanobot.config.paths import get_workspace_path, is_default_workspace
from nanobot.config.schema import Config
from nanobot.utils.helpers import sync_workspace_templates
from nanobot.utils.restart import (
consume_restart_notice_from_env,
format_restart_completed_message,
should_show_cli_restart_notice,
)
app = typer.Typer(
name="nanobot",
@ -206,57 +211,6 @@ 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).
@ -649,7 +603,6 @@ 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
@ -748,8 +701,6 @@ 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."""
@ -826,13 +777,6 @@ 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(),
@ -874,7 +818,6 @@ 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)
@ -915,13 +858,12 @@ 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)
restart_notice = consume_restart_notice_from_env()
if restart_notice and should_show_cli_restart_notice(restart_notice, session_id):
_print_agent_response(
format_restart_completed_message(restart_notice.started_at_raw),
render_markdown=False,
)
# Shared reference for progress callbacks
_thinking: ThinkingSpinner | None = None
@ -960,6 +902,11 @@ 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()

5
nanobot/command/builtin.py
View File

@ -9,8 +9,8 @@ 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
from nanobot.utils.restart import set_restart_notice_to_env
async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
@ -36,8 +36,7 @@ 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
set_restart_notice_to_env(channel=msg.channel, chat_id=msg.chat_id)
async def _do_restart():
await asyncio.sleep(1)

4
nanobot/config/runtime_keys.py
View File

@ -1,4 +0,0 @@
"""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"

58
nanobot/utils/restart.py Normal file
View File

@ -0,0 +1,58 @@
"""Helpers for restart notification messages."""
from __future__ import annotations
import os
import time
from dataclasses import dataclass
RESTART_NOTIFY_CHANNEL_ENV = "NANOBOT_RESTART_NOTIFY_CHANNEL"
RESTART_NOTIFY_CHAT_ID_ENV = "NANOBOT_RESTART_NOTIFY_CHAT_ID"
RESTART_STARTED_AT_ENV = "NANOBOT_RESTART_STARTED_AT"
@dataclass(frozen=True)
class RestartNotice:
channel: str
chat_id: str
started_at_raw: str
def format_restart_completed_message(started_at_raw: str) -> str:
"""Build restart completion text and include elapsed time when available."""
elapsed_suffix = ""
if started_at_raw:
try:
elapsed_s = max(0.0, time.time() - float(started_at_raw))
elapsed_suffix = f" in {elapsed_s:.1f}s"
except ValueError:
pass
return f"Restart completed{elapsed_suffix}."
def set_restart_notice_to_env(*, channel: str, chat_id: str) -> None:
"""Write restart notice env values for the next process."""
os.environ[RESTART_NOTIFY_CHANNEL_ENV] = channel
os.environ[RESTART_NOTIFY_CHAT_ID_ENV] = chat_id
os.environ[RESTART_STARTED_AT_ENV] = str(time.time())
def consume_restart_notice_from_env() -> RestartNotice | None:
"""Read and clear restart notice env values once for this process."""
channel = os.environ.pop(RESTART_NOTIFY_CHANNEL_ENV, "").strip()
chat_id = os.environ.pop(RESTART_NOTIFY_CHAT_ID_ENV, "").strip()
started_at_raw = os.environ.pop(RESTART_STARTED_AT_ENV, "").strip()
if not (channel and chat_id):
return None
return RestartNotice(channel=channel, chat_id=chat_id, started_at_raw=started_at_raw)
def should_show_cli_restart_notice(notice: RestartNotice, session_id: str) -> bool:
"""Return True when a restart notice should be shown in this CLI session."""
if notice.channel != "cli":
return False
if ":" in session_id:
_, cli_chat_id = session_id.split(":", 1)
else:
cli_chat_id = session_id
return not notice.chat_id or notice.chat_id == cli_chat_id

28
tests/channels/test_channel_plugins.py
View File

@ -13,6 +13,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.channels.manager import ChannelManager
from nanobot.config.schema import ChannelsConfig
from nanobot.utils.restart import RestartNotice
# ---------------------------------------------------------------------------
@ -929,3 +930,30 @@ async def test_start_all_creates_dispatch_task():
# Dispatch task should have been created
assert mgr._dispatch_task is not None
@pytest.mark.asyncio
async def test_notify_restart_done_enqueues_outbound_message():
"""Restart notice should schedule send_with_retry for target channel."""
fake_config = SimpleNamespace(
channels=ChannelsConfig(),
providers=SimpleNamespace(groq=SimpleNamespace(api_key="")),
)
mgr = ChannelManager.__new__(ChannelManager)
mgr.config = fake_config
mgr.bus = MessageBus()
mgr.channels = {"feishu": _StartableChannel(fake_config, mgr.bus)}
mgr._dispatch_task = None
mgr._send_with_retry = AsyncMock()
notice = RestartNotice(channel="feishu", chat_id="oc_123", started_at_raw="100.0")
with patch("nanobot.channels.manager.consume_restart_notice_from_env", return_value=notice):
mgr._notify_restart_done_if_needed()
await asyncio.sleep(0)
mgr._send_with_retry.assert_awaited_once()
sent_channel, sent_msg = mgr._send_with_retry.await_args.args
assert sent_channel is mgr.channels["feishu"]
assert sent_msg.channel == "feishu"
assert sent_msg.chat_id == "oc_123"
assert sent_msg.content.startswith("Restart completed")

81
tests/cli/test_restart_command.py
View File

@ -5,7 +5,6 @@ from __future__ import annotations
import asyncio
import os
import time
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -37,8 +36,12 @@ 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
from nanobot.utils.restart import (
RESTART_NOTIFY_CHANNEL_ENV,
RESTART_NOTIFY_CHAT_ID_ENV,
RESTART_STARTED_AT_ENV,
)
loop, bus = _make_loop()
msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart")
@ -50,6 +53,7 @@ class TestRestartCommand:
assert "Restarting" in out.content
assert os.environ.get(RESTART_NOTIFY_CHANNEL_ENV) == "cli"
assert os.environ.get(RESTART_NOTIFY_CHAT_ID_ENV) == "direct"
assert os.environ.get(RESTART_STARTED_AT_ENV)
await asyncio.sleep(1.5)
mock_execv.assert_called_once()
@ -196,76 +200,3 @@ 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

49
tests/utils/test_restart.py Normal file
View File

@ -0,0 +1,49 @@
"""Tests for restart notice helpers."""
from __future__ import annotations
import os
from nanobot.utils.restart import (
RestartNotice,
consume_restart_notice_from_env,
format_restart_completed_message,
set_restart_notice_to_env,
should_show_cli_restart_notice,
)
def test_set_and_consume_restart_notice_env_roundtrip(monkeypatch):
monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_CHANNEL", raising=False)
monkeypatch.delenv("NANOBOT_RESTART_NOTIFY_CHAT_ID", raising=False)
monkeypatch.delenv("NANOBOT_RESTART_STARTED_AT", raising=False)
set_restart_notice_to_env(channel="feishu", chat_id="oc_123")
notice = consume_restart_notice_from_env()
assert notice is not None
assert notice.channel == "feishu"
assert notice.chat_id == "oc_123"
assert notice.started_at_raw
# Consumed values should be cleared from env.
assert consume_restart_notice_from_env() is None
assert "NANOBOT_RESTART_NOTIFY_CHANNEL" not in os.environ
assert "NANOBOT_RESTART_NOTIFY_CHAT_ID" not in os.environ
assert "NANOBOT_RESTART_STARTED_AT" not in os.environ
def test_format_restart_completed_message_with_elapsed(monkeypatch):
monkeypatch.setattr("nanobot.utils.restart.time.time", lambda: 102.0)
assert format_restart_completed_message("100.0") == "Restart completed in 2.0s."
def test_should_show_cli_restart_notice():
notice = RestartNotice(channel="cli", chat_id="direct", started_at_raw="100")
assert should_show_cli_restart_notice(notice, "cli:direct") is True
assert should_show_cli_restart_notice(notice, "cli:other") is False
assert should_show_cli_restart_notice(notice, "direct") is True
non_cli = RestartNotice(channel="feishu", chat_id="oc_1", started_at_raw="100")
assert should_show_cli_restart_notice(non_cli, "cli:direct") is False