From ec8dee802c3727e6293e1d0bba9c6d0bb171b718 Mon Sep 17 00:00:00 2001
From: Re-bin <xubinrencs@gmail.com>
Date: Fri, 27 Feb 2026 02:39:38 +0000
Subject: [PATCH] refactor: simplify message tool suppress and inline
 consolidation locks

---
 README.md                        |  2 +-
 nanobot/agent/loop.py            | 41 ++++++++++----------------------
 tests/test_consolidate_offset.py |  2 +-
 3 files changed, 14 insertions(+), 31 deletions(-)

diff --git a/README.md b/README.md
index be360dc..71922fb 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@
 
 ⚡️ Delivers core agent functionality in just **~4,000** lines of code — **99% smaller** than Clawdbot's 430k+ lines.
 
-📏 Real-time line count: **3,966 lines** (run `bash core_agent_lines.sh` to verify anytime)
+📏 Real-time line count: **3,932 lines** (run `bash core_agent_lines.sh` to verify anytime)
 
 ## 📢 News
 
diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py
index 6155f99..e3a9d67 100644
--- a/nanobot/agent/loop.py
+++ b/nanobot/agent/loop.py
@@ -43,6 +43,8 @@ class AgentLoop:
     5. Sends responses back
     """
 
+    _TOOL_RESULT_MAX_CHARS = 500
+
     def __init__(
         self,
         bus: MessageBus,
@@ -145,17 +147,10 @@ class AgentLoop:
 
     def _set_tool_context(self, channel: str, chat_id: str, message_id: str | None = None) -> None:
         """Update context for all tools that need routing info."""
-        if message_tool := self.tools.get("message"):
-            if isinstance(message_tool, MessageTool):
-                message_tool.set_context(channel, chat_id, message_id)
-
-        if spawn_tool := self.tools.get("spawn"):
-            if isinstance(spawn_tool, SpawnTool):
-                spawn_tool.set_context(channel, chat_id)
-
-        if cron_tool := self.tools.get("cron"):
-            if isinstance(cron_tool, CronTool):
-                cron_tool.set_context(channel, chat_id)
+        for name in ("message", "spawn", "cron"):
+            if tool := self.tools.get(name):
+                if hasattr(tool, "set_context"):
+                    tool.set_context(channel, chat_id, *([message_id] if name == "message" else []))
 
     @staticmethod
     def _strip_think(text: str | None) -> str | None:
@@ -315,18 +310,6 @@ class AgentLoop:
         self._running = False
         logger.info("Agent loop stopping")
 
-    def _get_consolidation_lock(self, session_key: str) -> asyncio.Lock:
-        lock = self._consolidation_locks.get(session_key)
-        if lock is None:
-            lock = asyncio.Lock()
-            self._consolidation_locks[session_key] = lock
-        return lock
-
-    def _prune_consolidation_lock(self, session_key: str, lock: asyncio.Lock) -> None:
-        """Drop lock entry if no longer in use."""
-        if not lock.locked():
-            self._consolidation_locks.pop(session_key, None)
-
     async def _process_message(
         self,
         msg: InboundMessage,
@@ -362,7 +345,7 @@ class AgentLoop:
         # Slash commands
         cmd = msg.content.strip().lower()
         if cmd == "/new":
-            lock = self._get_consolidation_lock(session.key)
+            lock = self._consolidation_locks.setdefault(session.key, asyncio.Lock())
             self._consolidating.add(session.key)
             try:
                 async with lock:
@@ -383,7 +366,8 @@ class AgentLoop:
                 )
             finally:
                 self._consolidating.discard(session.key)
-                self._prune_consolidation_lock(session.key, lock)
+                if not lock.locked():
+                    self._consolidation_locks.pop(session.key, None)
 
             session.clear()
             self.sessions.save(session)
@@ -397,7 +381,7 @@ class AgentLoop:
         unconsolidated = len(session.messages) - session.last_consolidated
         if (unconsolidated >= self.memory_window and session.key not in self._consolidating):
             self._consolidating.add(session.key)
-            lock = self._get_consolidation_lock(session.key)
+            lock = self._consolidation_locks.setdefault(session.key, asyncio.Lock())
 
             async def _consolidate_and_unlock():
                 try:
@@ -405,7 +389,8 @@ class AgentLoop:
                         await self._consolidate_memory(session)
                 finally:
                     self._consolidating.discard(session.key)
-                    self._prune_consolidation_lock(session.key, lock)
+                    if not lock.locked():
+                        self._consolidation_locks.pop(session.key, None)
                     _task = asyncio.current_task()
                     if _task is not None:
                         self._consolidation_tasks.discard(_task)
@@ -454,8 +439,6 @@ class AgentLoop:
             metadata=msg.metadata or {},
         )
 
-    _TOOL_RESULT_MAX_CHARS = 500
-
     def _save_turn(self, session: Session, messages: list[dict], skip: int) -> None:
         """Save new-turn messages into session, truncating large tool results."""
         from datetime import datetime
diff --git a/tests/test_consolidate_offset.py b/tests/test_consolidate_offset.py
index 323519e..6755124 100644
--- a/tests/test_consolidate_offset.py
+++ b/tests/test_consolidate_offset.py
@@ -812,7 +812,7 @@ class TestConsolidationDeduplicationGuard:
         loop.sessions.save(session)
 
         # Ensure lock exists before /new.
-        _ = loop._get_consolidation_lock(session.key)
+        loop._consolidation_locks.setdefault(session.key, asyncio.Lock())
         assert session.key in loop._consolidation_locks
 
         async def _ok_consolidate(sess, archive_all: bool = False) -> bool: