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
59
nanobot/agent/commands.py
Normal file
59
nanobot/agent/commands.py
Normal file
@@ -0,0 +1,59 @@
|
||||
"""Command definitions and dispatch for the agent loop.
|
||||
|
||||
Commands are slash-prefixed messages (e.g. /stop, /new, /help) that are
|
||||
handled specially — either immediately in the run() loop or inside
|
||||
_process_message before the LLM is called.
|
||||
|
||||
To add a new command:
|
||||
1. Add a CommandDef to COMMANDS
|
||||
2. If immediate=True, add a handler in AgentLoop._handle_immediate_command
|
||||
3. If immediate=False, add handling in AgentLoop._process_message
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CommandDef:
|
||||
"""Definition of a slash command."""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
immediate: bool = False # True = handled in run() loop, bypasses message processing
|
||||
|
||||
|
||||
# Registry of all known commands.
|
||||
# "immediate" commands are handled while the agent may be busy (e.g. /stop).
|
||||
# Non-immediate commands go through normal _process_message flow.
|
||||
COMMANDS: dict[str, CommandDef] = {
|
||||
"/stop": CommandDef("/stop", "Stop the current task", immediate=True),
|
||||
"/new": CommandDef("/new", "Start a new conversation"),
|
||||
"/help": CommandDef("/help", "Show available commands"),
|
||||
}
|
||||
|
||||
|
||||
def parse_command(text: str) -> str | None:
|
||||
"""Extract a slash command from message text.
|
||||
|
||||
Returns the command string (e.g. "/stop") or None if not a command.
|
||||
"""
|
||||
stripped = text.strip()
|
||||
if not stripped.startswith("/"):
|
||||
return None
|
||||
return stripped.split()[0].lower()
|
||||
|
||||
|
||||
def is_immediate_command(cmd: str) -> bool:
|
||||
"""Check if a command should be handled immediately, bypassing processing."""
|
||||
defn = COMMANDS.get(cmd)
|
||||
return defn.immediate if defn else False
|
||||
|
||||
|
||||
def get_help_text() -> str:
|
||||
"""Generate help text from registered commands."""
|
||||
lines = ["🐈 nanobot commands:"]
|
||||
for defn in COMMANDS.values():
|
||||
lines.append(f"{defn.name} — {defn.description}")
|
||||
return "\n".join(lines)
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user