feat: dimos spy - universal transport spy over LCM + Zenoh

dimos/protocol/pubsub/impl/zenohpubsub.py

@@ -197,61 +197,31 @@ def unsubscribe() -> None:
         return unsubscribe
 
     def subscribe_all(self, callback: Callable[[bytes, Topic], Any]) -> Callable[[], None]:
-        """Subscribe to all dimos topics, delivering only the latest per topic.
+        """Subscribe to every key, delivering every message (non-conflating).
 
-        Unlike `subscribe`, this is best effort. If it's done otherwise, rerun lags behind.
+        "All topics" means everything this session can observe ('**'). Consumers
+        that only want the newest message per topic use subscribe_latest().
         """
-        latest: dict[str, tuple[bytes, Topic]] = {}
-        lock = threading.Lock()
-        wake = threading.Event()
-        stop = threading.Event()
-
-        def collect(msg: bytes, topic: Topic) -> None:
-            # Fast path on the Zenoh delivery thread: keep only the newest per topic.
-            with lock:
-                latest[str(topic)] = (msg, topic)
-            wake.set()
-
-        def drain() -> None:
-            while not stop.is_set():
-                wake.wait()
-                wake.clear()
-                with lock:
-                    batch = list(latest.values())
-                    latest.clear()
-                for msg, topic in batch:
-                    try:
-                        callback(msg, topic)
-                    except Exception:
-                        logger.error("Error in subscribe_all callback", exc_info=True)
-
-        thread = threading.Thread(target=drain, name="zenoh-subscribe-all", daemon=True)
-
-        def stop_drain() -> None:
-            stop.set()
-            wake.set()  # unblock the drain so it observes the stop flag
-            thread.join(timeout=2.0)
-
-        # Register the stop callback and launch the drain thread under the lock so
-        # stop() can't run between them and miss the thread. If stop() already ran,
-        # bail without starting anything.
+        return self.subscribe(Topic("**"), callback)
+
+    def _register_drain_stop(
+        self, stop_drain: Callable[[], None], thread: threading.Thread
+    ) -> bool:
+        # Register + start under the subscriber lock so stop() can't run between
+        # them and miss the thread. If stop() already ran, bail without starting.
         with self._subscriber_lock:
             if self._stopped:
-                return lambda: None
+                return False
             self._drain_stops.append(stop_drain)
             thread.start()
+        return True
 
-        inner_unsub = self.subscribe(Topic("dimos/**"), collect)
-
-        def unsubscribe() -> None:
-            with self._subscriber_lock:
-                if stop_drain not in self._drain_stops:
-                    return  # Already removed by stop() or a concurrent unsubscribe
-                self._drain_stops.remove(stop_drain)
-            inner_unsub()
-            stop_drain()
-
-        return unsubscribe
+    def _unregister_drain_stop(self, stop_drain: Callable[[], None]) -> bool:
+        with self._subscriber_lock:
+            if stop_drain not in self._drain_stops:
+                return False  # Already removed by stop() or a concurrent unsubscribe
+            self._drain_stops.remove(stop_drain)
+        return True
 
     def stop(self) -> None:
         with self._subscriber_lock:

dimos/protocol/pubsub/spec.py

@@ -17,8 +17,13 @@
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+import threading
 from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
 
+from dimos.utils.logging_config import setup_logger
+
+logger = setup_logger()
+
 MsgT = TypeVar("MsgT")
 TopicT = TypeVar("TopicT")
 MsgT_co = TypeVar("MsgT_co", covariant=True)
@@ -115,7 +120,100 @@ def subscribe(
 #
 # - DiscoveryPubSub: Native support for discovering new topics as they appear.
 #   Provides a default subscribe_all() by subscribing to each discovered topic.
-class AllPubSub(PubSub[TopicT, MsgT], ABC):
+class SubscribeLatestMixin(Generic[TopicT, MsgT], ABC):
+    """Conflated subscribe-all, built on top of subscribe_all().
+
+    Mixed into AllPubSub and DiscoveryPubSub so both expose subscribe_latest().
+    """
+
+    @abstractmethod
+    def subscribe_all(self, callback: Callable[[MsgT, TopicT], Any]) -> Callable[[], None]: ...
+
+    def _register_drain_stop(
+        self, stop_drain: Callable[[], None], thread: threading.Thread
+    ) -> bool:
+        """Start ``thread`` and record ``stop_drain`` for transport-level shutdown.
+
+        Default: just start the thread. Transports whose stop() must join live
+        drains (e.g. zenoh) override this to register under their own lock and
+        return False if already stopped (subscribe_latest then bails).
+        """
+        thread.start()
+        return True
+
+    def _unregister_drain_stop(self, stop_drain: Callable[[], None]) -> bool:
+        """Forget a drain stopper. Returns False if already removed."""
+        return True
+
+    def subscribe_latest(self, callback: Callable[[MsgT, TopicT], Any]) -> Callable[[], None]:
+        """Subscribe to all topics, conflated: newest message per topic wins.
+
+        For consumers that only need the current value per topic and must not
+        lag behind a fast producer (e.g. the rerun bridge). When the callback
+        is slower than the message flow, intermediate messages on a topic are
+        dropped and only the latest is delivered.
+
+        Contract:
+        - The callback runs on a dedicated drain thread, never on the
+          transport's delivery thread.
+        - Per topic, the newest pending message is always eventually delivered
+          (no starvation); intermediate ones may be skipped.
+        - The returned callable unsubscribes the underlying subscribe_all and
+          stops+joins the drain thread.
+        """
+        latest: dict[str, tuple[MsgT, TopicT]] = {}
+        lock = threading.Lock()
+        wake = threading.Event()
+        stop = threading.Event()
+
+        def collect(msg: MsgT, topic: TopicT) -> None:
+            # Fast path on the transport's delivery thread: keep only newest per topic.
+            with lock:
+                latest[str(topic)] = (msg, topic)
+            wake.set()
+
+        def drain() -> None:
+            while not stop.is_set():
+                wake.wait()
+                wake.clear()
+                with lock:
+                    batch = list(latest.values())
+                    latest.clear()
+                for msg, topic in batch:
+                    try:
+                        callback(msg, topic)
+                    except Exception:
+                        logger.error("Error in subscribe_latest callback", exc_info=True)
+
+        thread = threading.Thread(target=drain, name="subscribe-latest-drain", daemon=True)
+
+        def stop_drain() -> None:
+            stop.set()
+            wake.set()  # unblock the drain so it observes the stop flag
+            thread.join(timeout=2.0)
+
+        # Register + start the drain atomically w.r.t. transport shutdown, then
+        # subscribe. If the transport already stopped, bail without subscribing.
+        if not self._register_drain_stop(stop_drain, thread):
+            return lambda: None
+        try:
+            inner_unsub = self.subscribe_all(collect)
+        except BaseException:
+            # Don't leak the drain thread if the underlying subscribe fails.
+            if self._unregister_drain_stop(stop_drain):
+                stop_drain()
+            raise
+
+        def unsubscribe() -> None:
+            if not self._unregister_drain_stop(stop_drain):
+                return  # Already removed by stop() or a concurrent unsubscribe
+            inner_unsub()
+            stop_drain()
+
+        return unsubscribe
+
+
+class AllPubSub(SubscribeLatestMixin[TopicT, MsgT], PubSub[TopicT, MsgT], ABC):
     """Mixin for PubSub that supports subscribing to all topics.
 
     Subclass from this if you support native subscribe-all (e.g. MQTT #, Redis *).
@@ -124,7 +222,13 @@ class AllPubSub(PubSub[TopicT, MsgT], ABC):
 
     @abstractmethod
     def subscribe_all(self, callback: Callable[[MsgT, TopicT], Any]) -> Callable[[], None]:
-        """Subscribe to all topics."""
+        """Subscribe to all topics.
+
+        Contract: delivers EVERY message on every topic this transport can
+        observe — implementations must not conflate, sample, or drop beyond
+        the transport's own delivery semantics. Consumers that only want the
+        newest message per topic use subscribe_latest() instead.
+        """
         ...
 
     def subscribe_new_topics(self, callback: Callable[[TopicT], Any]) -> Callable[[], None]:
@@ -144,7 +248,7 @@ def on_msg(msg: MsgT, topic: TopicT) -> None:
 
 
 # This is for ros for now
-class DiscoveryPubSub(PubSub[TopicT, MsgT], ABC):
+class DiscoveryPubSub(SubscribeLatestMixin[TopicT, MsgT], PubSub[TopicT, MsgT], ABC):
     """Mixin for PubSub that supports discovery of topics.
 
     Subclass from this if you support topic discovery (e.g. MQTT, Redis, NATS, RabbitMQ).
@@ -187,5 +291,9 @@ class SubscribeAllCapable(Protocol[MsgT_co, TopicT_co]):
     """
 
     def subscribe_all(self, callback: Callable[[Any, Any], Any]) -> Callable[[], None]:
-        """Subscribe to all messages on all topics."""
+        """Subscribe to all messages on all topics (every message, no conflation)."""
+        ...
+
+    def subscribe_latest(self, callback: Callable[[Any, Any], Any]) -> Callable[[], None]:
+        """Subscribe to all topics, conflated to the newest message per topic."""
         ...