feat: extensible command system + task-based dispatch with /stop

- Add commands.py with CommandDef registry, parse_command(), get_help_text()
- Refactor run() to dispatch messages as asyncio tasks (non-blocking)
- /stop is an 'immediate' command: handled inline, cancels active task
- Global processing lock serializes message handling (safe for shared state)
- _pending_tasks set prevents GC of dispatched tasks before lock acquisition
- _dispatch() registers/clears active tasks, catches CancelledError gracefully
- /help now auto-generated from COMMANDS registry

Closes #849
This commit is contained in:
coldxiangyu
2026-02-25 17:51:00 +08:00
parent 9e806d7159
commit 3c12efa728
3 changed files with 349 additions and 17 deletions
+70 -17
View File
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable
from loguru import logger
from nanobot.agent.commands import get_help_text, is_immediate_command, parse_command
from nanobot.agent.context import ContextBuilder
from nanobot.agent.memory import MemoryStore
from nanobot.agent.subagent import SubagentManager
@@ -99,6 +100,9 @@ class AgentLoop:
self._consolidating: set[str] = set() # Session keys with consolidation in progress
self._consolidation_tasks: set[asyncio.Task] = set() # Strong refs to in-flight tasks
self._consolidation_locks: dict[str, asyncio.Lock] = {}
self._active_tasks: dict[str, asyncio.Task] = {} # session_key -> running task
self._pending_tasks: set[asyncio.Task] = set() # Strong refs until dispatch starts
self._processing_lock = asyncio.Lock() # Serialize message processing
self._register_default_tools()
def _register_default_tools(self) -> None:
@@ -238,7 +242,12 @@ class AgentLoop:
return final_content, tools_used, messages
async def run(self) -> None:
"""Run the agent loop, processing messages from the bus."""
"""Run the agent loop, processing messages from the bus.
Regular messages are dispatched as asyncio tasks so the loop stays
responsive to immediate commands like /stop. A global processing
lock serializes message handling to avoid shared-state races.
"""
self._running = True
await self._connect_mcp()
logger.info("Agent loop started")
@@ -249,24 +258,68 @@ class AgentLoop:
self.bus.consume_inbound(),
timeout=1.0
)
try:
response = await self._process_message(msg)
if response is not None:
await self.bus.publish_outbound(response)
elif msg.channel == "cli":
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, content="", metadata=msg.metadata or {},
))
except Exception as e:
logger.error("Error processing message: {}", e)
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel,
chat_id=msg.chat_id,
content=f"Sorry, I encountered an error: {str(e)}"
))
# Immediate commands (/stop) are handled inline
cmd = parse_command(msg.content)
if cmd and is_immediate_command(cmd):
await self._handle_immediate_command(cmd, msg)
continue
# Regular messages (including non-immediate commands) are
# dispatched as tasks so the loop keeps consuming.
task = asyncio.create_task(self._dispatch(msg))
self._pending_tasks.add(task)
task.add_done_callback(self._pending_tasks.discard)
except asyncio.TimeoutError:
continue
async def _handle_immediate_command(self, cmd: str, msg: InboundMessage) -> None:
"""Handle a command that must be processed while the agent may be busy."""
if cmd == "/stop":
task = self._active_tasks.get(msg.session_key)
if task and not task.done():
task.cancel()
try:
await task
except (asyncio.CancelledError, Exception):
pass
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="⏹ Task stopped.",
))
else:
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="No active task to stop.",
))
async def _dispatch(self, msg: InboundMessage) -> None:
"""Dispatch a message for processing under the global lock."""
async with self._processing_lock:
self._active_tasks[msg.session_key] = asyncio.current_task() # type: ignore[arg-type]
try:
response = await self._process_message(msg)
if response is not None:
await self.bus.publish_outbound(response)
elif msg.channel == "cli":
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="", metadata=msg.metadata or {},
))
except asyncio.CancelledError:
logger.info("Task cancelled for session {}", msg.session_key)
# Response already sent by _handle_immediate_command
except Exception as e:
logger.error("Error processing message: {}", e)
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel,
chat_id=msg.chat_id,
content=f"Sorry, I encountered an error: {str(e)}"
))
finally:
self._active_tasks.pop(msg.session_key, None)
async def close_mcp(self) -> None:
"""Close MCP connections."""
if self._mcp_stack:
@@ -358,7 +411,7 @@ class AgentLoop:
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/help — Show available commands")
content=get_help_text())
unconsolidated = len(session.messages) - session.last_consolidated
if (unconsolidated >= self.memory_window and session.key not in self._consolidating):