Hua/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: background post-response memory consolidation for faster replies

Browse Source
This commit is contained in:
Xubin Ren
2026-03-16 09:01:11 +00:00
parent 6d63e22e86
commit 46b19b15e1
5 changed files with 23 additions and 522 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
nanobot/agent/loop.py
View File

@@ -100,7 +100,7 @@ class AgentLoop:
self._mcp_connected = False
self._mcp_connecting = False
self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks
self._pending_archives: list[asyncio.Task] = []
self._background_tasks: list[asyncio.Task] = []
self._processing_lock = asyncio.Lock()
self.memory_consolidator = MemoryConsolidator(
workspace=workspace,
@@ -257,8 +257,6 @@ class AgentLoop:
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
self._running = True
await self._connect_mcp()
# Start background consolidation task
await self.memory_consolidator.start_background_task()
logger.info("Agent loop started")
while self._running:
@@ -334,9 +332,9 @@ class AgentLoop:
async def close_mcp(self) -> None:
"""Drain pending background archives, then close MCP connections."""
if self._pending_archives:
await asyncio.gather(*self._pending_archives, return_exceptions=True)
self._pending_archives.clear()
if self._background_tasks:
await asyncio.gather(*self._background_tasks, return_exceptions=True)
self._background_tasks.clear()
if self._mcp_stack:
try:
await self._mcp_stack.aclose()
@@ -344,11 +342,16 @@ class AgentLoop:
pass # MCP SDK cancel scope cleanup is noisy but harmless
self._mcp_stack = None
async def stop(self) -> None:
"""Stop the agent loop and background tasks."""
def _schedule_background(self, coro) -> None:
"""Schedule a coroutine as a tracked background task (drained on shutdown)."""
task = asyncio.create_task(coro)
self._background_tasks.append(task)
task.add_done_callback(self._background_tasks.remove)
def stop(self) -> None:
"""Stop the agent loop."""
self._running = False
await self.memory_consolidator.stop_background_task()
logger.info("Agent loop stopped")
logger.info("Agent loop stopping")
async def _process_message(
self,
@@ -364,8 +367,7 @@ class AgentLoop:
logger.info("Processing system message from {}", msg.sender_id)
key = f"{channel}:{chat_id}"
session = self.sessions.get_or_create(key)
self.memory_consolidator.record_activity(key)
await self.memory_consolidator.maybe_consolidate_by_tokens_async(session)
await self.memory_consolidator.maybe_consolidate_by_tokens(session)
self._set_tool_context(channel, chat_id, msg.metadata.get("message_id"))
history = session.get_history(max_messages=0)
messages = self.context.build_messages(
@@ -375,6 +377,7 @@ class AgentLoop:
final_content, _, all_msgs = await self._run_agent_loop(messages)
self._save_turn(session, all_msgs, 1 + len(history))
self.sessions.save(session)
self._schedule_background(self.memory_consolidator.maybe_consolidate_by_tokens(session))
return OutboundMessage(channel=channel, chat_id=chat_id,
content=final_content or "Background task completed.")
@@ -383,7 +386,6 @@ class AgentLoop:
key = session_key or msg.session_key
session = self.sessions.get_or_create(key)
self.memory_consolidator.record_activity(key)
# Slash commands
cmd = msg.content.strip().lower()
@@ -394,11 +396,7 @@ class AgentLoop:
self.sessions.invalidate(session.key)
if snapshot:
task = asyncio.create_task(
self.memory_consolidator.archive_messages(snapshot)
)
self._pending_archives.append(task)
task.add_done_callback(self._pending_archives.remove)
self._schedule_background(self.memory_consolidator.archive_messages(snapshot))
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
content="New session started.")
@@ -413,8 +411,7 @@ class AgentLoop:
return OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, content="\n".join(lines),
)
# Record activity and schedule background consolidation for non-slash commands
self.memory_consolidator.record_activity(key)
await self.memory_consolidator.maybe_consolidate_by_tokens(session)
self._set_tool_context(msg.channel, msg.chat_id, msg.metadata.get("message_id"))
if message_tool := self.tools.get("message"):
@@ -446,6 +443,7 @@ class AgentLoop:
self._save_turn(session, all_msgs, 1 + len(history))
self.sessions.save(session)
self._schedule_background(self.memory_consolidator.maybe_consolidate_by_tokens(session))
if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn:
return None