Update node_transport.py

node_transport.py

@@ -2,34 +2,6 @@
 node_transport.py
 =================
 Transport-neutral node bus for edge fleet coordination.
-
-Semantics (all backends):
-  - Latest-value, non-consuming reads
-  - Soft TTL (STATE_TTL_SEC); recv returns None when stale
-  - No pickle / no eval — versioned binary or JSON only
-
-Wire envelope (both hidden state and product alerts):
-  [version:u8][flags:u8][payload...]
-  flags bit0 = zlib-compressed payload
-
-Hidden inner payload (unchanged for torch side):
-  [n_layers:u32][batch:u32][hidden:u32][float32 data...]  little-endian
-
-Alert inner payload:
-  UTF-8 JSON object (see alert_to_bytes / bytes_to_alert)
-
-NodeTransport protocol (Genius-safe abstraction):
-  send/recv_hidden_state, send/recv_alert, clear
-
-MQTT is one backend (hub-and-spoke). A future libp2p / SuperGenius backend
-implements the same protocol without callers knowing topic names or brokers.
-
-Steps covered here:
-  1. Product alerts on the protocol (not MQTT-only helpers)
-  2. MQTT hardening: TLS, auth, LWT + birth presence, reconnect
-  3. MQTT 5 when available (message expiry, user properties); 3.1.1 fallback
-  4. Schema version + optional zlib on all wire payloads
-  5. Multi-broker list (try in order) — topology as config, not hard-coded sites
 """
 
 from __future__ import annotations
@@ -64,11 +36,7 @@ STATE_TTL_SEC    = 300                # soft TTL for latest-value reads
 WIRE_VERSION     = 1
 FLAG_COMPRESSED  = 0x01
 
-# MQTT topic layout is an implementation detail of MQTTTransport only.
 DEFAULT_TOPIC_ROOT = "weather"
-# weather/hidden/{zone_id}
-# weather/alert/{zone_id}
-# weather/presence/{node_id}
 
 
 # ---------------------------------------------------------------------------
@@ -76,7 +44,6 @@ DEFAULT_TOPIC_ROOT = "weather"
 # ---------------------------------------------------------------------------
 
 def pack_wire(payload: bytes, *, compress: bool = False) -> bytes:
-    """Prefix payload with version + flags. Optional zlib on payload only."""
     flags = 0
     body = payload
     if compress and len(payload) > 64:
@@ -86,15 +53,9 @@ def pack_wire(payload: bytes, *, compress: bool = False) -> bytes:
 
 
 def unpack_wire(data: bytes) -> bytes:
-    """
-    Strip version envelope. Accepts:
-      - versioned: [ver][flags][payload]
-      - legacy hidden: raw [u32][u32][u32][floats...] (no version byte)
-    """
     if len(data) < 2:
         raise ValueError(f"wire payload too short: {len(data)}")
 
-    # Legacy hidden: first uint32 is n_layers in 1..8 and total length matches
     if _looks_like_legacy_hidden(data):
         return data
 
@@ -125,15 +86,6 @@ def _looks_like_legacy_hidden(data: bytes) -> bool:
 # ---------------------------------------------------------------------------
 
 def hidden_to_bytes(hidden_tensor) -> bytes:
-    """
-    Serialise a GRU hidden state tensor to inner payload bytes.
-
-    Format:
-        [n_layers: uint32][batch_size: uint32][hidden_size: uint32][float32 data...]
-    Always little-endian, always float32.
-    Call pack_wire(...) before sending on the bus if compression/versioning needed;
-    LocalTransport and MQTTTransport pack automatically on send.
-    """
     import torch
 
     h = hidden_tensor.detach().cpu().float()
@@ -144,10 +96,6 @@ def hidden_to_bytes(hidden_tensor) -> bytes:
 
 
 def bytes_to_hidden(data: bytes):
-    """
-    Safe deserialisation with strict validation before any ML import.
-    Accepts wire envelope or legacy/inner payload.
-    """
     if _looks_like_legacy_hidden(data):
         inner = data
     else:
@@ -188,11 +136,6 @@ def bytes_to_hidden(data: bytes):
 
 @dataclass
 class ProductAlert:
-    """Client/product-facing alert snapshot for the node bus.
-
-    Mirrors env terminate info + scorer freeze fields. Transport-agnostic:
-    any NodeTransport backend can carry the same bytes.
-    """
     zone_id: str
     product_actionable: bool
     elevated: bool
@@ -212,7 +155,6 @@ class ProductAlert:
 
 
 def alert_to_bytes(alert: ProductAlert) -> bytes:
-    """JSON inner payload for product alerts (UTF-8)."""
     raw = json.dumps(alert.to_dict(), separators=(",", ":"), ensure_ascii=True).encode("utf-8")
     if len(raw) > MAX_ALERT_BYTES:
         raise ValueError(f"alert payload exceeds MAX_ALERT_BYTES ({len(raw)})")
@@ -220,12 +162,10 @@ def alert_to_bytes(alert: ProductAlert) -> bytes:
 
 
 def bytes_to_alert(data: bytes) -> ProductAlert:
-    """Parse alert from wire or inner JSON bytes."""
     try:
         inner = unpack_wire(data)
     except ValueError:
         inner = data
-    # If unpack produced garbage for pure JSON, try raw
     try:
         obj = json.loads(inner.decode("utf-8"))
     except (UnicodeDecodeError, json.JSONDecodeError):
@@ -266,7 +206,6 @@ def product_alert_from_risk(
     trigger: str = "",
     **extra: Any,
 ) -> ProductAlert:
-    """Helper for env / edge callers building an alert from scorer/env info."""
     return ProductAlert(
         zone_id=zone_id,
         product_actionable=product_actionable,
@@ -287,12 +226,6 @@ def product_alert_from_risk(
 
 @runtime_checkable
 class NodeTransport(Protocol):
-    """Minimal interface every transport must implement.
-
-    Hidden state and product alerts are peer channels with identical
-    latest-value / TTL semantics. Topic names, brokers, and P2P routes
-    must not leak into callers.
-    """
 
     async def send_hidden_state(self, zone_id: str, hidden: bytes) -> None:
         ...
@@ -315,14 +248,6 @@ class NodeTransport(Protocol):
 # ---------------------------------------------------------------------------
 
 class LocalTransport:
-    """
-    In-process deterministic transport.
-
-    Semantics:
-      - Latest-value (NOT consuming)
-      - TTL enforced on recv
-      - Auto wire-envelope on send when payload is bare inner
-    """
 
     def __init__(self, *, compress: bool = False, ttl_sec: float = STATE_TTL_SEC) -> None:
         self._hidden: Dict[str, Tuple[bytes, float]] = {}
@@ -331,7 +256,6 @@ class LocalTransport:
         self._ttl = float(ttl_sec)
 
     def _pack(self, payload: bytes) -> bytes:
-        # If already versioned, leave as-is
         if len(payload) >= 2 and payload[0] == WIRE_VERSION and not _looks_like_legacy_hidden(payload):
             return payload
         return pack_wire(payload, compress=self._compress)
@@ -380,23 +304,6 @@ class LocalTransport:
 # ---------------------------------------------------------------------------
 
 class MQTTTransport:
-    """
-    MQTT backend for NodeTransport (hub-and-spoke).
-
-    Guarantees:
-      - Same latest-value / non-consuming / TTL semantics as LocalTransport
-      - Thread-safe receive buffers
-      - Fallback to LocalTransport when broker unavailable
-      - TLS + username/password when configured
-      - LWT + retained birth on presence topic
-      - MQTT 5 message expiry when broker/protocol supports it; 3.1.1 fallback
-      - Multi-broker: try hosts in order (step 5 topology as config)
-
-    Topic layout (implementation detail — not part of NodeTransport):
-      {topic_root}/hidden/{zone_id}
-      {topic_root}/alert/{zone_id}
-      {topic_root}/presence/{node_id}
-    """
 
     def __init__(
         self,
@@ -405,21 +312,16 @@ class MQTTTransport:
         username: Optional[str] = None,
         password: Optional[str] = None,
         topic_root: str = DEFAULT_TOPIC_ROOT,
-        # legacy alias
         topic_prefix: Optional[str] = None,
         *,
-        # step 5: try multiple brokers in order
         brokers: Optional[Sequence[Tuple[str, int]]] = None,
-        # step 2: TLS
         tls: bool = False,
         tls_ca_certs: Optional[str] = None,
         tls_certfile: Optional[str] = None,
         tls_keyfile: Optional[str] = None,
         tls_insecure: bool = False,
-        # identity / presence
         client_id: Optional[str] = None,
         node_id: str = "node-0",
-        # behaviour
         compress: bool = False,
         ttl_sec: float = STATE_TTL_SEC,
         keepalive: int = 60,
@@ -434,9 +336,7 @@ class MQTTTransport:
         self.username = username
         self.password = password
         self.topic_root = topic_root.rstrip("/")
-        # backward compat: old topic_prefix meant hidden prefix
         if topic_prefix is not None:
-            # if caller passed weather/hidden, derive root
             parts = topic_prefix.rstrip("/").split("/")
             if parts and parts[-1] == "hidden":
                 self.topic_root = "/".join(parts[:-1]) or DEFAULT_TOPIC_ROOT
@@ -516,7 +416,6 @@ class MQTTTransport:
         if self.client_id:
             client_kwargs["client_id"] = self.client_id
 
-        # Prefer MQTT 5 when requested and available
         self._mqttv5_active = False
         if self.use_mqttv5 and hasattr(mqtt, "MQTTv5"):
             protocol = mqtt.MQTTv5
@@ -524,7 +423,6 @@ class MQTTTransport:
 
         client_kwargs["protocol"] = protocol
         if callback_api is not None:
-            # paho 2.x
             try:
                 client = mqtt.Client(
                     callback_api_version=callback_api.VERSION2,
@@ -550,7 +448,6 @@ class MQTTTransport:
             if self.tls_insecure:
                 client.tls_insecure_set(True)
 
-        # LWT: retained offline presence
         lwt_payload = json.dumps(
             {"node_id": self.node_id, "status": "offline", "ts": time.time()},
             separators=(",", ":"),
@@ -566,7 +463,6 @@ class MQTTTransport:
         client.on_disconnect = self._on_disconnect
         client.on_message = self._on_message
 
-        # Automatic reconnect (paho)
         try:
             client.reconnect_delay_set(min_delay=1, max_delay=30)
         except Exception:
@@ -578,7 +474,6 @@ class MQTTTransport:
         client.loop_start()
 
         self._client = client
-        # Brief wait for on_connect — non-blocking best-effort
         for _ in range(20):
             if self._connected:
                 break
@@ -588,7 +483,6 @@ class MQTTTransport:
             "MQTTTransport: connect attempted %s:%s mqttv5=%s connected=%s",
             host, port, self._mqttv5_active, self._connected,
         )
-        # Consider success if loop started; on_connect may race
         return True
 
     def _teardown_client(self) -> None:
@@ -602,7 +496,6 @@ class MQTTTransport:
         self._connected = False
 
     def _on_connect(self, client, userdata, flags, reason_code, properties=None) -> None:
-        # Compatible with both paho 1.x (rc int) and 2.x VERSION2 (reason_code)
         rc = reason_code
         if hasattr(reason_code, "value"):
             rc = reason_code.value
@@ -615,7 +508,6 @@ class MQTTTransport:
         if not self._connected:
             return
 
-        # Retained birth / online presence
         birth = json.dumps(
             {"node_id": self.node_id, "status": "online", "ts": time.time()},
             separators=(",", ":"),
@@ -630,7 +522,6 @@ class MQTTTransport:
         except Exception as e:
             logger.warning("MQTTTransport: birth publish failed: %s", e)
 
-        # Re-subscribe after reconnect
         try:
             client.subscribe(f"{self.topic_root}/hidden/#", qos=self.qos)
             client.subscribe(f"{self.topic_root}/alert/#", qos=self.qos)
@@ -638,7 +529,6 @@ class MQTTTransport:
             logger.warning("MQTTTransport: resubscribe failed: %s", e)
 
     def _on_disconnect(self, client, userdata, flags, reason_code=None, properties=None) -> None:
-        # paho 1.x: (client, userdata, rc); 2.x VERSION2 adds flags/properties
         self._connected = False
         logger.warning("MQTTTransport: disconnected reason=%s", reason_code)
 
@@ -688,9 +578,7 @@ class MQTTTransport:
                 from paho.mqtt.packettypes import PacketTypes
 
                 props = Properties(PacketTypes.PUBLISH)
-                # Message expiry (seconds) — MQTT 5 equivalent of soft TTL
                 props.MessageExpiryInterval = int(self._ttl)
-                # User properties for schema / debugging (not required by receivers)
                 try:
                     props.UserProperty = [("wire_v", str(WIRE_VERSION))]
                 except Exception:
@@ -726,7 +614,6 @@ class MQTTTransport:
         with self._lock:
             item = self._hidden_rx.get(zone_id)
         if item is None:
-            # also check fallback (local writes while disconnected)
             return await self._fallback.recv_hidden_state(zone_id)
 
         data, ts = item
@@ -774,7 +661,6 @@ class MQTTTransport:
 
     async def disconnect(self) -> None:
         if self._client is not None:
-            # Best-effort offline presence before disconnect
             try:
                 offline = json.dumps(
                     {"node_id": self.node_id, "status": "offline", "ts": time.time()},
@@ -800,14 +686,6 @@ def create_node_transport(
     use_mqtt: bool = False,
     **mqtt_kwargs: Any,
 ) -> NodeTransport:
-    """
-    Public factory.
-
-    Usage:
-        training → create_node_transport()
-        edge     → create_node_transport(use_mqtt=True, broker="...", tls=True, ...)
-        multi    → create_node_transport(use_mqtt=True, brokers=[("a", 8883), ("b", 8883)], tls=True)
-    """
     if use_mqtt:
         return MQTTTransport(**mqtt_kwargs)
     return LocalTransport(
@@ -826,7 +704,6 @@ def _self_test() -> None:
     async def _run() -> None:
         print("node_transport.py self-test")
 
-        # --- wire pack/unpack ---
         raw = b"hello-alert-payload"
         wire = pack_wire(raw, compress=False)
         assert unpack_wire(wire) == raw
@@ -834,8 +711,6 @@ def _self_test() -> None:
         assert unpack_wire(wire_c) == raw * 20
         print("  wire envelope OK")
 
-        # --- hidden round-trip (no torch required if we skip tensor path) ---
-        # Build a synthetic inner payload
         n_layers, batch, hidden = 1, 1, 8
         inner = struct.pack("<III", n_layers, batch, hidden) + (b"\x00\x00\x00\x00" * (n_layers * batch * hidden))
         assert _looks_like_legacy_hidden(inner)
@@ -843,7 +718,6 @@ def _self_test() -> None:
         assert unpack_wire(w) == inner
         print("  hidden wire OK")
 
-        # --- alert serialisation ---
         alert = product_alert_from_risk(
             "karawang_rice",
             alert_level="warning",
@@ -862,7 +736,6 @@ def _self_test() -> None:
         assert abs(alert2.drought_risk - 0.42) < 1e-6
         print("  alert serialisation OK")
 
-        # --- LocalTransport protocol ---
         tr = LocalTransport()
         await tr.send_hidden_state("z1", inner)
         got = await tr.recv_hidden_state("z1")
@@ -880,9 +753,7 @@ def _self_test() -> None:
         assert await tr.recv_alert("z1") is None
         print("  LocalTransport protocol OK")
 
-        # --- MQTTTransport falls back without broker ---
         mqtt_tr = MQTTTransport(broker="127.0.0.1", port=1, node_id="test-node")
-        # do not connect — should use fallback
         await mqtt_tr.send_alert("z2", ab)
         got_b = await mqtt_tr.recv_alert("z2")
         assert got_b is not None
@@ -891,7 +762,6 @@ def _self_test() -> None:
         assert await mqtt_tr.recv_hidden_state("z2") is not None
         print("  MQTTTransport offline fallback OK")
 
-        # --- factory ---
         t0 = create_node_transport()
         assert isinstance(t0, LocalTransport)
         t1 = create_node_transport(use_mqtt=True, broker="localhost")
@@ -905,4 +775,4 @@ def _self_test() -> None:
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
-    _self_test()
+    _self_test()