From c865b293a9a772c000eed0cda63eecbd2efa472e Mon Sep 17 00:00:00 2001
From: Darye <54469750+DaryeDev@users.noreply.github.com>
Date: Wed, 18 Feb 2026 20:18:27 +0100
Subject: [PATCH] feat: enhance message context handling by adding message_id
 parameter

---
 nanobot/agent/loop.py          |  8 ++++----
 nanobot/agent/tools/message.py | 14 +++++++++++---
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py
index e5a5183..7855297 100644
--- a/nanobot/agent/loop.py
+++ b/nanobot/agent/loop.py
@@ -133,11 +133,11 @@ class AgentLoop:
         await self._mcp_stack.__aenter__()
         await connect_mcp_servers(self._mcp_servers, self.tools, self._mcp_stack)
 
-    def _set_tool_context(self, channel: str, chat_id: str) -> None:
+    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_tool.set_context(channel, chat_id, message_id)
 
         if spawn_tool := self.tools.get("spawn"):
             if isinstance(spawn_tool, SpawnTool):
@@ -321,7 +321,7 @@ class AgentLoop:
         if len(session.messages) > self.memory_window:
             asyncio.create_task(self._consolidate_memory(session))
 
-        self._set_tool_context(msg.channel, msg.chat_id)
+        self._set_tool_context(msg.channel, msg.chat_id, msg.metadata.get("message_id"))
         initial_messages = self.context.build_messages(
             history=session.get_history(max_messages=self.memory_window),
             current_message=msg.content,
@@ -379,7 +379,7 @@ class AgentLoop:
         
         session_key = f"{origin_channel}:{origin_chat_id}"
         session = self.sessions.get_or_create(session_key)
-        self._set_tool_context(origin_channel, origin_chat_id)
+        self._set_tool_context(origin_channel, origin_chat_id, msg.metadata.get("message_id"))
         initial_messages = self.context.build_messages(
             history=session.get_history(max_messages=self.memory_window),
             current_message=msg.content,
diff --git a/nanobot/agent/tools/message.py b/nanobot/agent/tools/message.py
index 3853725..10947c4 100644
--- a/nanobot/agent/tools/message.py
+++ b/nanobot/agent/tools/message.py
@@ -13,16 +13,19 @@ class MessageTool(Tool):
         self, 
         send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
         default_channel: str = "",
-        default_chat_id: str = ""
+        default_chat_id: str = "",
+        default_message_id: str | None = None
     ):
         self._send_callback = send_callback
         self._default_channel = default_channel
         self._default_chat_id = default_chat_id
+        self._default_message_id = default_message_id
     
-    def set_context(self, channel: str, chat_id: str) -> None:
+    def set_context(self, channel: str, chat_id: str, message_id: str | None = None) -> None:
         """Set the current message context."""
         self._default_channel = channel
         self._default_chat_id = chat_id
+        self._default_message_id = message_id
     
     def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None:
         """Set the callback for sending messages."""
@@ -67,11 +70,13 @@ class MessageTool(Tool):
         content: str, 
         channel: str | None = None, 
         chat_id: str | None = None,
+        message_id: str | None = None,
         media: list[str] | None = None,
         **kwargs: Any
     ) -> str:
         channel = channel or self._default_channel
         chat_id = chat_id or self._default_chat_id
+        message_id = message_id or self._default_message_id
         
         if not channel or not chat_id:
             return "Error: No target channel/chat specified"
@@ -83,7 +88,10 @@ class MessageTool(Tool):
             channel=channel,
             chat_id=chat_id,
             content=content,
-            media=media or []
+            media=media or [],
+            metadata={
+                "message_id": message_id,
+            }
         )
         
         try: