fix(restart): send completion notice after channel is ready and unify runtime keys

This commit is contained in:
imfondof 2026-04-02 16:42:47 +08:00 committed by Xubin Ren
parent a05f83da89
commit ba7c07ccf2
4 changed files with 156 additions and 6 deletions

View File

@ -206,6 +206,57 @@ def _is_exit_command(command: str) -> bool:
return command.lower() in EXIT_COMMANDS 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: async def _read_interactive_input_async() -> str:
"""Read user input using prompt_toolkit (handles paste, history, display). """Read user input using prompt_toolkit (handles paste, history, display).
@ -598,6 +649,7 @@ def gateway(
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.channels.manager import ChannelManager 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.service import CronService
from nanobot.cron.types import CronJob from nanobot.cron.types import CronJob
from nanobot.heartbeat.service import HeartbeatService from nanobot.heartbeat.service import HeartbeatService
@ -696,6 +748,8 @@ def gateway(
# Create channel manager # Create channel manager
channels = ChannelManager(config, bus) 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]: def _pick_heartbeat_target() -> tuple[str, str]:
"""Pick a routable channel/chat target for heartbeat-triggered messages.""" """Pick a routable channel/chat target for heartbeat-triggered messages."""
@ -772,6 +826,13 @@ def gateway(
try: try:
await cron.start() await cron.start()
await heartbeat.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( await asyncio.gather(
agent.run(), agent.run(),
channels.start_all(), channels.start_all(),
@ -813,6 +874,7 @@ def agent(
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus 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 from nanobot.cron.service import CronService
config = _load_runtime_config(config, workspace) config = _load_runtime_config(config, workspace)
@ -853,6 +915,13 @@ def agent(
channels_config=config.channels, channels_config=config.channels,
timezone=config.agents.defaults.timezone, 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 # Shared reference for progress callbacks
_thinking: ThinkingSpinner | None = None _thinking: ThinkingSpinner | None = None
@ -891,11 +960,6 @@ def agent(
_init_prompt_session() _init_prompt_session()
console.print(f"{__logo__} Interactive mode (type [bold]exit[/bold] or [bold]Ctrl+C[/bold] to quit)\n") 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): def _handle_signal(signum, frame):
sig_name = signal.Signals(signum).name sig_name = signal.Signals(signum).name
_restore_terminal() _restore_terminal()

View File

@ -9,6 +9,7 @@ import sys
from nanobot import __version__ from nanobot import __version__
from nanobot.bus.events import OutboundMessage from nanobot.bus.events import OutboundMessage
from nanobot.command.router import CommandContext, CommandRouter 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.helpers import build_status_content
@ -35,6 +36,8 @@ async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
async def cmd_restart(ctx: CommandContext) -> OutboundMessage: async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
"""Restart the process in-place via os.execv.""" """Restart the process in-place via os.execv."""
msg = ctx.msg 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(): async def _do_restart():
await asyncio.sleep(1) await asyncio.sleep(1)

View File

@ -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"

View File

@ -3,7 +3,9 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import os
import time import time
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
import pytest import pytest
@ -35,15 +37,19 @@ class TestRestartCommand:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_restart_sends_message_and_calls_execv(self): async def test_restart_sends_message_and_calls_execv(self):
from nanobot.command.builtin import cmd_restart 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.command.router import CommandContext
loop, bus = _make_loop() loop, bus = _make_loop()
msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart") 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) 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) out = await cmd_restart(ctx)
assert "Restarting" in out.content 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) await asyncio.sleep(1.5)
mock_execv.assert_called_once() mock_execv.assert_called_once()
@ -190,3 +196,76 @@ class TestRestartCommand:
assert response is not None assert response is not None
assert response.metadata == {"render_as": "text"} 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