From 64aeeceed02aadb19e51f82d71674024baec4b95 Mon Sep 17 00:00:00 2001 From: Re-bin Date: Thu, 12 Mar 2026 04:33:51 +0000 Subject: [PATCH] Add /restart command: restart the bot process from any channel --- nanobot/agent/loop.py | 43 +++++++++++++------- tests/test_restart_command.py | 76 +++++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 15 deletions(-) create mode 100644 tests/test_restart_command.py diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 597f852..5fe0ee0 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -4,8 +4,8 @@ from __future__ import annotations import asyncio import json -import re import os +import re import sys from contextlib import AsyncExitStack from pathlib import Path @@ -258,8 +258,11 @@ class AgentLoop: except asyncio.TimeoutError: continue - if msg.content.strip().lower() == "/stop": + cmd = msg.content.strip().lower() + if cmd == "/stop": await self._handle_stop(msg) + elif cmd == "/restart": + await self._handle_restart(msg) else: task = asyncio.create_task(self._dispatch(msg)) self._active_tasks.setdefault(msg.session_key, []).append(task) @@ -276,11 +279,23 @@ class AgentLoop: pass sub_cancelled = await self.subagents.cancel_by_session(msg.session_key) total = cancelled + sub_cancelled - content = f"⏹ Stopped {total} task(s)." if total else "No active task to stop." + content = f"Stopped {total} task(s)." if total else "No active task to stop." await self.bus.publish_outbound(OutboundMessage( channel=msg.channel, chat_id=msg.chat_id, content=content, )) + async def _handle_restart(self, msg: InboundMessage) -> None: + """Restart the process in-place via os.execv.""" + await self.bus.publish_outbound(OutboundMessage( + channel=msg.channel, chat_id=msg.chat_id, content="Restarting...", + )) + + async def _do_restart(): + await asyncio.sleep(1) + os.execv(sys.executable, [sys.executable] + sys.argv) + + asyncio.create_task(_do_restart()) + async def _dispatch(self, msg: InboundMessage) -> None: """Process a message under the global lock.""" async with self._processing_lock: @@ -375,18 +390,16 @@ class AgentLoop: return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, content="New session started.") if cmd == "/help": - return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, - content="🐈 nanobot commands:\n/new — Start a new conversation\n/stop — Stop the current task\n/help — Show available commands") - if cmd == "/restart": - await self.bus.publish_outbound(OutboundMessage( - channel=msg.channel, chat_id=msg.chat_id, content="🔄 Restarting..." - )) - async def _r(): - await asyncio.sleep(1) - os.execv(sys.executable, [sys.executable] + sys.argv) - asyncio.create_task(_r()) - return None - + lines = [ + "🐈 nanobot commands:", + "/new — Start a new conversation", + "/stop — Stop the current task", + "/restart — Restart the bot", + "/help — Show available commands", + ] + return OutboundMessage( + channel=msg.channel, chat_id=msg.chat_id, content="\n".join(lines), + ) await self.memory_consolidator.maybe_consolidate_by_tokens(session) self._set_tool_context(msg.channel, msg.chat_id, msg.metadata.get("message_id")) diff --git a/tests/test_restart_command.py b/tests/test_restart_command.py new file mode 100644 index 0000000..c495347 --- /dev/null +++ b/tests/test_restart_command.py @@ -0,0 +1,76 @@ +"""Tests for /restart slash command.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +from nanobot.bus.events import InboundMessage + + +def _make_loop(): + """Create a minimal AgentLoop with mocked dependencies.""" + from nanobot.agent.loop import AgentLoop + from nanobot.bus.queue import MessageBus + + bus = MessageBus() + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + workspace = MagicMock() + workspace.__truediv__ = MagicMock(return_value=MagicMock()) + + with patch("nanobot.agent.loop.ContextBuilder"), \ + patch("nanobot.agent.loop.SessionManager"), \ + patch("nanobot.agent.loop.SubagentManager"): + loop = AgentLoop(bus=bus, provider=provider, workspace=workspace) + return loop, bus + + +class TestRestartCommand: + + @pytest.mark.asyncio + async def test_restart_sends_message_and_calls_execv(self): + loop, bus = _make_loop() + msg = InboundMessage(channel="cli", sender_id="user", chat_id="direct", content="/restart") + + with patch("nanobot.agent.loop.os.execv") as mock_execv: + await loop._handle_restart(msg) + out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0) + assert "Restarting" in out.content + + await asyncio.sleep(1.5) + mock_execv.assert_called_once() + + @pytest.mark.asyncio + async def test_restart_intercepted_in_run_loop(self): + """Verify /restart is handled at the run-loop level, not inside _dispatch.""" + loop, bus = _make_loop() + msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/restart") + + with patch.object(loop, "_handle_restart") as mock_handle: + mock_handle.return_value = None + await bus.publish_inbound(msg) + + loop._running = True + run_task = asyncio.create_task(loop.run()) + await asyncio.sleep(0.1) + loop._running = False + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + mock_handle.assert_called_once() + + @pytest.mark.asyncio + async def test_help_includes_restart(self): + loop, bus = _make_loop() + msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/help") + + response = await loop._process_message(msg) + + assert response is not None + assert "/restart" in response.content