re-wrote message bus to remove return value and add run()

agentserver/message_bus.py

@@ -1,6 +1,10 @@
+# agentserver/message_bus.py
+# Refactored January 01, 2026 – MessageBus with run() pump and out-of-band shutdown
+
 import asyncio
 import logging
-from typing import Dict, Optional, Callable
+from typing import AsyncIterator, Callable, Dict, Optional, Awaitable
+
 from lxml import etree
 
 from .xml_listener import XMLListener
@@ -15,21 +19,26 @@ logger = logging.getLogger("agentserver.bus")
 
 
 class MessageBus:
-    """
+    """The sovereign message carrier.
-    The Pure Carrier: Routes XML trees between sovereign listeners.
+
-    - Core: High-speed tree-to-tree routing.
+    - Routes canonical XML trees by root tag and <to/> meta.
-    - Border: 'Air Lock' for ingesting raw bytes.
+    - Pure dispatch: tree → optional response tree.
-    - Physics: Hardcoded hooks for the Audit Witness.
+    - Active pump via run(): handles serialization and egress.
+    - Out-of-band shutdown via asyncio.Event (fast-path, flood-immune).
     """
 
-    def __init__(self, log_hook: Callable[[etree.Element], None]):
+    def __init__(self, log_hook: Callable[[etree._Element], None]):
         # root_tag -> {agent_name -> XMLListener}
         self.listeners: Dict[str, Dict[str, XMLListener]] = {}
         # Global lookup for directed <to/> routing
         self.global_names: Dict[str, XMLListener] = {}
-        # The 'Physics' hook provided by the Host (AgentServer)
+
+        # The Sovereign Witness hook
         self._log_hook = log_hook
 
+        # Out-of-band shutdown signal (set only by AgentServer on privileged command)
+        self.shutdown_event = asyncio.Event()
+
     async def register_listener(self, listener: XMLListener) -> None:
         """Register an organ. Enforces global identity uniqueness."""
         if listener.agent_name in self.global_names:
@@ -42,85 +51,108 @@ class MessageBus:
 
         logger.info(f"Registered organ: {listener.agent_name}")
 
-
+    async def deliver_bytes(self, raw_xml: bytes, client_id: Optional[str] = None) -> None:
-    async def deliver_bytes(
+        """Air Lock: ingest raw bytes, repair/canonicalize, inject into core."""
-            self,
-            raw_xml: bytes,
-            client_id: Optional[str] = None
-    ) -> Optional[str]:
-        """
-        The Air Lock: Ingests raw bytes from unauthenticated or foreign sources.
-        Repairs and validates them before injecting the tree into the core.
-        """
         try:
-            # 1. Pressurize (Bytes -> Tree)
+            envelope_tree = repair_and_canonicalize(raw_xml)
-            envelope_tree: etree._Element = repair_and_canonicalize(raw_xml)
+            await self.dispatch(envelope_tree, client_id)
-            # 2. Inject into core
-            return await self.dispatch(envelope_tree, client_id)
         except XmlTamperError as e:
             logger.warning(f"Air Lock Reject: {e}")
-            return None
-
 
     async def dispatch(
         self,
         envelope_tree: etree._Element,
-            client_id: Optional[str] = None
+        client_id: Optional[str] = None,
-    ) -> Optional[str]:
+    ) -> etree._Element | None:
-        """
+        """Pure routing heart. Returns validated response tree or None."""
-        The Pure Heart: Routes a tree. Returns a validated response string.
+        # 1. WITNESS – every canonical envelope is seen
-        """
-        # 1. WITNESS: Automatic logging of the truth
         self._log_hook(envelope_tree)
 
-        # 2. Extract Metadata
+        # 2. Extract envelope metadata
         meta = envelope_tree.find(f"{ENV}meta")
+        if meta is None:
+            return None
         from_name = meta.findtext(f"{ENV}from")
         to_name = meta.findtext(f"{ENV}to")
-        thread_id = meta.findtext(f"{ENV}thread")
+        thread_id = meta.findtext(f"{ENV}thread_id") or meta.findtext(f"{ENV}thread")
 
-        payload_elem = next((c for c in envelope_tree.iterchildren() if c.tag != f"{ENV}meta"), None)
+        # Find payload (first non-meta child)
-        if payload_elem is None: return None
+        payload_elem = next((c for c in envelope_tree if c.tag != f"{ENV}meta"), None)
+        if payload_elem is None:
+            return None
         payload_tag = payload_elem.tag
 
-        # 3. AUTONOMIC REFLEX: The Logger Hook (Hardcoded Physics)
+        # 3. AUTONOMIC REFLEX: explicit <log/>
         if payload_tag == LOG_TAG:
-            self._log_hook(envelope_tree)
+            self._log_hook(envelope_tree)  # extra vent
-            return f"<message xmlns='{ENV_NS}'><meta><from>system</from><to>{from_name}</to><thread_id>{thread_id}</thread_id></meta><logged status='success'/></message>"
+            # Minimal ack envelope
+            ack = etree.Element(f"{ENV}message")
+            meta_ack = etree.SubElement(ack, f"{ENV}meta")
+            etree.SubElement(meta_ack, f"{ENV}from").text = "system"
+            if from_name:
+                etree.SubElement(meta_ack, f"{ENV}to").text = from_name
+            if thread_id:
+                etree.SubElement(meta_ack, f"{ENV}thread_id").text = thread_id
+            etree.SubElement(ack, "logged", status="success")
+            return ack
 
         # 4. ROUTING
         listeners_for_tag = self.listeners.get(payload_tag, {})
-        response_tree = None
+        response_tree: Optional[etree._Element] = None
         responding_agent_name = "unknown"
 
         if to_name:
-            # Directed Mode
+            # Directed
             target = listeners_for_tag.get(to_name) or self.global_names.get(to_name)
             if target:
                 responding_agent_name = target.agent_name
                 response_tree = await target.handle(envelope_tree, thread_id, from_name or client_id)
         else:
-            # Broadcast Mode
+            # Broadcast – first non-None wins (current policy)
-            tasks = [l.handle(envelope_tree, thread_id, from_name or client_id) for l in listeners_for_tag.values()]
+            tasks = [
+                l.handle(envelope_tree, thread_id, from_name or client_id)
+                for l in listeners_for_tag.values()
+            ]
             results = await asyncio.gather(*tasks, return_exceptions=True)
-            for i, resp in enumerate(results):
+            for listener, result in zip(listeners_for_tag.values(), results):
-                if isinstance(resp, etree._Element):
+                if isinstance(result, etree._Element):
-                    responding_agent_name = list(listeners_for_tag.values())[i].agent_name
+                    responding_agent_name = listener.agent_name
-                    response_tree = resp
+                    response_tree = result
+                    break  # first-wins
+
+        # 5. IDENTITY INSPECTION – prevent spoofing
+        if response_tree is not None:
+            actual_from = response_tree.findtext(f"{ENV}meta/{ENV}from")
+            if actual_from != responding_agent_name:
+                logger.critical(
+                    f"IDENTITY THEFT BLOCKED: expected {responding_agent_name}, got {actual_from}"
+                )
+                return None
+
+        return response_tree
+
+    async def run(
+        self,
+        inbound: AsyncIterator[etree._Element],
+        outbound: Callable[[bytes], Awaitable[None]],
+        client_id: Optional[str] = None,
+    ) -> None:
+        """Active pump for a connection. Handles serialization and egress."""
+        try:
+            async for envelope_tree in inbound:
+                if self.shutdown_event.is_set():
                     break
 
-        # 5. EGRESS CUSTOMS: Final validation before bytes hit the wire
+                response_tree = await self.dispatch(envelope_tree, client_id)
                 if response_tree is not None:
-            return self._inspect_and_serialize(response_tree, responding_agent_name)
+                    serialized = etree.tostring(
-
+                        response_tree, encoding="utf-8", pretty_print=True
-        return None
+                    )
-
+                    await outbound(serialized)
-
+        finally:
-    def _inspect_and_serialize(self, tree: etree._Element, expected_from: str) -> Optional[str]:
+            # Optional final courtesy message on clean exit
-        """Ensures identity integrity and performs final string serialization."""
+            goodbye = b"<message xmlns='https://xml-pipeline.org/ns/envelope/1'><goodbye reason='connection-closed'/></message>"
-        actual_from = tree.findtext(f"{ENV}meta/{ENV}from")
+            try:
-        if actual_from != expected_from:
+                await outbound(goodbye)
-            logger.critical(f"IDENTITY THEFT BLOCKED: {expected_from} spoofed {actual_from}")
+            except Exception:
-            return None
+                pass  # connection already gone
-        # noinspection PyTypeChecker
-        return etree.tostring(tree, encoding="unicode", pretty_print=True)