diff --git a/dimos/protocol/pubsub/impl/zenohpubsub.py b/dimos/protocol/pubsub/impl/zenohpubsub.py index afaa413c43..466eff6e1a 100644 --- a/dimos/protocol/pubsub/impl/zenohpubsub.py +++ b/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: diff --git a/dimos/protocol/pubsub/spec.py b/dimos/protocol/pubsub/spec.py index fe979fce82..1652621417 100644 --- a/dimos/protocol/pubsub/spec.py +++ b/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,94 @@ 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 + inner_unsub = self.subscribe_all(collect) + + 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 +216,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 +242,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 +285,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.""" ... diff --git a/dimos/protocol/pubsub/spy.py b/dimos/protocol/pubsub/spy.py new file mode 100644 index 0000000000..1e6086630f --- /dev/null +++ b/dimos/protocol/pubsub/spy.py @@ -0,0 +1,275 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transport-agnostic pubsub spy: topic discovery, rates, sizes, liveness. + +Design doc: docs/usage/transports/spy.md. Task spec: TASK.md (branch agent/spy-architect). + +HARD CONSTRAINT: the spy never decodes message payloads. Sources tap the +raw-bytes pubsub layer (LCMPubSubBase, ZenohPubSubBase — beneath the encoder +mixins), so the hot path per message is (topic string, payload length, +timestamp) and nothing else. Message *types* are still visible because they +are embedded in the topic string ("/cmd_vel#geometry_msgs.Twist"). + +Decoding is a separate, per-topic, opt-in concern: SpySource.subscribe_decoded +is the spec'd hook, not implemented in v1. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +import threading +import time +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + # The entire hot-path event: (topic string incl. '#type' suffix, wire payload length). + TapCallback = Callable[[str, int], None] + + +@dataclass(frozen=True, slots=True) +class SpyKey: + """Identity of one spied topic: which transport saw it + its raw topic string.""" + + transport: str # SpySource.name, e.g. "lcm", "zenoh" + topic: str # raw transport topic, e.g. "/cmd_vel#geometry_msgs.Twist" + + +def split_type_suffix(topic: str) -> tuple[str, str | None]: + """Split a spied topic string into (base_topic, msg_type_name or None). + + Both LCM and zenoh sources deliver topics in the uniform str(Topic) form + "base#pkg.Msg" (zenoh keys are converted back by _key_expr_to_topic). + Render-time helper — never called on the hot path. + + >>> split_type_suffix("/cmd_vel#geometry_msgs.Twist") + ('/cmd_vel', 'geometry_msgs.Twist') + >>> split_type_suffix("/plain") + ('/plain', None) + """ + base, sep, suffix = topic.partition("#") + return (base, suffix) if sep else (topic, None) + + +class TopicStats: + """Sliding-window traffic statistics for one spied topic. + + Records only (timestamp, nbytes) pairs — no payloads are retained. + Timestamps are passed in explicitly (callers use time.time(); tests inject + values), so all stats are deterministic functions of recorded data. + + Thread-safety: record() may be called from transport threads while readers + query from a UI thread. + """ + + total_bytes: int + total_msgs: int + last_seen: float | None # timestamp of newest recorded message, None if none yet + + def __init__(self, history_window: float = 60.0) -> None: + """history_window: seconds of per-message history kept for windowed stats. + + total_bytes/total_msgs/last_seen survive eviction; only windowed stats + (freq/bytes_per_sec/avg_size) forget evicted messages. + """ + self.history_window = history_window + self._history: deque[tuple[float, int]] = deque() # (timestamp, nbytes) + self._lock = threading.Lock() + self.total_bytes = 0 + self.total_msgs = 0 + self.last_seen = None + + def record(self, nbytes: int, timestamp: float) -> None: + """Hot path: O(1) amortized append + eviction of entries older than history_window.""" + with self._lock: + self._history.append((timestamp, nbytes)) + self.total_bytes += nbytes + self.total_msgs += 1 + self.last_seen = timestamp + cutoff = timestamp - self.history_window + while self._history and self._history[0][0] < cutoff: + self._history.popleft() + + def _in_window(self, window: float, now: float) -> list[tuple[float, int]]: + cutoff = now - window + with self._lock: + return [(ts, n) for ts, n in self._history if ts >= cutoff] + + def freq(self, window: float, now: float) -> float: + """Messages per second over [now - window, now]. 0.0 if none.""" + msgs = self._in_window(window, now) + return len(msgs) / window if msgs else 0.0 + + def bytes_per_sec(self, window: float, now: float) -> float: + """Payload bytes per second over [now - window, now]. 0.0 if none.""" + msgs = self._in_window(window, now) + return sum(n for _, n in msgs) / window if msgs else 0.0 + + def avg_size(self, window: float, now: float) -> float: + """Mean payload size in bytes over [now - window, now]. 0.0 if none.""" + msgs = self._in_window(window, now) + return sum(n for _, n in msgs) / len(msgs) if msgs else 0.0 + + +@runtime_checkable +class SpySource(Protocol): + """One transport's raw firehose feeding the spy. + + Invariants for implementations: + - tap() delivers EVERY observable message exactly once per callback + (non-conflating; built on the raw-bytes bus's subscribe_all). + - The tap callback receives (topic_str, nbytes) where topic_str is the + uniform str(Topic) form and nbytes is the wire payload length. + - Never decodes payloads, never retains them past the callback. + - start() is required before tap() delivers; stop() releases the bus. + """ + + name: str + + def start(self) -> None: ... + + def stop(self) -> None: ... + + def tap(self, callback: Callable[[str, int], None]) -> Callable[[], None]: ... + + def subscribe_decoded( + self, topic: str, callback: Callable[[Any], None] + ) -> Callable[[], None]: ... + + +class LCMSpySource: + """Spy source over LCM, via LCMPubSubBase.subscribe_all (raw regex '.*'). + + Delivers raw channel strings incl. the '#pkg.Msg' suffix. Payloads are the + LCM-encoded bytes; nbytes = len(payload). + """ + + name = "lcm" + + def __init__(self, **lcm_kwargs: Any) -> None: + """lcm_kwargs forwarded to LCMPubSubBase (e.g. lcm_url).""" + from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase + + self._bus = LCMPubSubBase(**lcm_kwargs) + + def start(self) -> None: + self._bus.start() + + def stop(self) -> None: + self._bus.stop() + + def tap(self, callback: Callable[[str, int], None]) -> Callable[[], None]: + return self._bus.subscribe_all(lambda msg, topic: callback(str(topic), len(msg))) + + def subscribe_decoded(self, topic: str, callback: Callable[[Any], None]) -> Callable[[], None]: + """Opt-in per-topic decoded tap — OFF the spy hot path. Not implemented in v1.""" + raise NotImplementedError + + +class ZenohSpySource: + """Spy source over zenoh, via ZenohPubSubBase.subscribe_all (all keys, '**'). + + Depends on the subscribe_all non-conflation fix (see TASK.md): every sample + must reach the tap. nbytes = payload length; topics arrive as str(Topic) + with the type suffix reconstructed from the key expression. + """ + + name = "zenoh" + + def __init__(self, **zenoh_kwargs: Any) -> None: + """zenoh_kwargs forwarded to ZenohPubSubBase (e.g. mode/connect/listen, session_pool).""" + from dimos.protocol.pubsub.impl.zenohpubsub import ZenohPubSubBase + + self._bus = ZenohPubSubBase(**zenoh_kwargs) + + def start(self) -> None: + self._bus.start() + + def stop(self) -> None: + self._bus.stop() + + def tap(self, callback: Callable[[str, int], None]) -> Callable[[], None]: + return self._bus.subscribe_all(lambda msg, topic: callback(str(topic), len(msg))) + + def subscribe_decoded(self, topic: str, callback: Callable[[Any], None]) -> Callable[[], None]: + """Opt-in per-topic decoded tap — OFF the spy hot path. Not implemented in v1.""" + raise NotImplementedError + + +class TransportSpy: + """Aggregates SpySources into per-(transport, topic) stats plus global totals. + + Owns the sources' lifecycle: start() starts every source and taps it; + stop() untaps and stops them. Stats rows appear lazily as topics are first + seen (a topic with no traffic since start is invisible — the spy observes, + it does not enumerate). + + Thread-safety: tap callbacks arrive on transport threads; snapshot() may be + called from any thread and returns a consistent view for rendering. + """ + + totals: TopicStats # all messages across all sources + + def __init__( + self, sources: Sequence[SpySource] | None = None, history_window: float = 60.0 + ) -> None: + """sources=None means default_sources().""" + self._sources = list(sources) if sources is not None else default_sources() + self._history_window = history_window + self._stats: dict[SpyKey, TopicStats] = {} + self._lock = threading.Lock() + self._untaps: list[Callable[[], None]] = [] + self.totals = TopicStats(history_window=history_window) + + def _tap_callback(self, transport: str) -> Callable[[str, int], None]: + def on_message(topic: str, nbytes: int) -> None: + now = time.time() + key = SpyKey(transport, topic) + with self._lock: + stats = self._stats.get(key) + if stats is None: + stats = TopicStats(history_window=self._history_window) + self._stats[key] = stats + stats.record(nbytes, now) + self.totals.record(nbytes, now) + + return on_message + + def start(self) -> None: + for source in self._sources: + source.start() + self._untaps.append(source.tap(self._tap_callback(source.name))) + + def stop(self) -> None: + for untap in self._untaps: + untap() + self._untaps.clear() + for source in self._sources: + source.stop() + + def snapshot(self) -> dict[SpyKey, TopicStats]: + """Current per-topic stats, safe to iterate while messages keep arriving.""" + with self._lock: + return dict(self._stats) + + +def default_sources() -> list[SpySource]: + """The spy observes ALL transports simultaneously, regardless of DIMOS_TRANSPORT. + + v1: [LCMSpySource(), ZenohSpySource()]. SHM/ROS/DDS/Redis are future sources. + """ + return [LCMSpySource(), ZenohSpySource()] diff --git a/dimos/protocol/pubsub/test_spy.py b/dimos/protocol/pubsub/test_spy.py new file mode 100644 index 0000000000..a61f5198aa --- /dev/null +++ b/dimos/protocol/pubsub/test_spy.py @@ -0,0 +1,447 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contract tests for `dimos spy` (TASK.md on branch agent/spy-architect). + +These encode the acceptance criteria and FAIL until the task is implemented: +- TopicStats: deterministic windowed stats from injected timestamps. +- subscribe_all: delivers EVERY message on LCM and zenoh (non-conflating). +- subscribe_latest: conflates to the newest message per topic. +- SpySources: count every message without ever decoding a payload. +- TransportSpy: merges sources into (transport, topic)-keyed stats + totals. +- subscribe_decoded: spec'd hook, must stay NotImplementedError in v1. +""" + +from collections.abc import Callable +from contextlib import contextmanager +import threading +import time +from typing import Any + +import pytest + +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic +from dimos.protocol.pubsub.impl.zenohpubsub import ZenohPubSubBase +from dimos.protocol.pubsub.spec import AllPubSub +from dimos.protocol.pubsub.spy import ( + LCMSpySource, + SpyKey, + TopicStats, + TransportSpy, + ZenohSpySource, + split_type_suffix, +) +from dimos.protocol.service.zenohservice import ZenohSessionPool + +VEC = Vector3(1.0, 2.0, 3.0) +VEC_BYTES = VEC.lcm_encode() + + +# TopicStats: pure, deterministic (injected timestamps, no sleeps) + + +def test_topic_stats_counts_and_totals(): + s = TopicStats(history_window=60.0) + t0 = 1000.0 + for i in range(10): + s.record(100, t0 + i * 0.1) + assert s.total_msgs == 10 + assert s.total_bytes == 1000 + assert s.last_seen == pytest.approx(t0 + 0.9) + + +def test_topic_stats_windowed_rates(): + s = TopicStats(history_window=60.0) + t0 = 1000.0 + for i in range(10): # 1 Hz, 50 bytes each + s.record(50, t0 + i) + now = t0 + 9.0 + assert s.freq(5.0, now) == pytest.approx(1.0, rel=0.25) + assert s.bytes_per_sec(5.0, now) == pytest.approx(50.0, rel=0.25) + assert s.avg_size(5.0, now) == pytest.approx(50.0) + + +def test_topic_stats_empty(): + s = TopicStats() + assert s.total_msgs == 0 + assert s.total_bytes == 0 + assert s.last_seen is None + assert s.freq(5.0, now=123.0) == 0.0 + assert s.bytes_per_sec(5.0, now=123.0) == 0.0 + assert s.avg_size(5.0, now=123.0) == 0.0 + + +def test_topic_stats_history_eviction_keeps_totals(): + s = TopicStats(history_window=60.0) + s.record(100, 0.0) + s.record(100, 1000.0) # first record is now far outside the history window + assert s.freq(60.0, now=1000.0) == pytest.approx(1 / 60.0) + assert s.total_msgs == 2 # totals survive eviction + assert s.total_bytes == 200 + + +# split_type_suffix: render-time topic parsing + + +def test_split_type_suffix(): + assert split_type_suffix("/cmd_vel#geometry_msgs.Twist") == ("/cmd_vel", "geometry_msgs.Twist") + assert split_type_suffix("dimos/cmd_vel#geometry_msgs.Twist") == ( + "dimos/cmd_vel", + "geometry_msgs.Twist", + ) + assert split_type_suffix("/plain") == ("/plain", None) + + +# subscribe_all contract: EVERY message, no conflation (spec-level fix) + + +@contextmanager +def lcm_base_ctx(): + bus = LCMPubSubBase() + bus.start() + try: + yield bus, Topic("/spy_contract", Vector3) + finally: + bus.stop() + + +@contextmanager +def zenoh_base_ctx(): + pool = ZenohSessionPool() + bus = ZenohPubSubBase(session_pool=pool) + bus.start() + try: + yield bus, Topic("dimos/spy_contract", Vector3) + finally: + bus.stop() + pool.close_all() + + +@pytest.mark.parametrize("bus_ctx", [lcm_base_ctx, zenoh_base_ctx], ids=["lcm", "zenoh"]) +def test_subscribe_all_delivers_every_message(bus_ctx): + """A gated (slow) consumer must still receive every message eventually. + + This is the sharp edge of the non-conflation contract: zenoh's original + subscribe_all conflated to latest-per-topic and fails this with ~2 of 50. + """ + n = 50 + with bus_ctx() as (bus, topic): + gate = threading.Event() + got = [] + done = threading.Event() + + def cb(msg, t): + gate.wait(15.0) # simulate a consumer slower than the burst + got.append((str(t), len(msg))) + if len(got) >= n: + done.set() + + unsub = bus.subscribe_all(cb) + time.sleep(0.5) # let the wildcard subscription establish + + for _ in range(n): + bus.publish(topic, VEC_BYTES) + gate.set() + + done.wait(10.0) + ours = [g for g in got if g[0] == str(topic)] + assert len(ours) == n + assert all(size == len(VEC_BYTES) for _, size in ours) + unsub() + + +# subscribe_latest contract: conflation is the explicit opt-in + + +class FakeAllPubSub(AllPubSub[str, bytes]): + """In-memory AllPubSub: publish() synchronously fans out to subscribers.""" + + def __init__(self): + self._subs: list[Callable[[bytes, str], None]] = [] + + def publish(self, topic: str, message: bytes) -> None: + for cb in list(self._subs): + cb(message, topic) + + def subscribe(self, topic, callback): + return self.subscribe_all(callback) + + def subscribe_all(self, callback): + self._subs.append(callback) + + def unsub(): + if callback in self._subs: + self._subs.remove(callback) + + return unsub + + +def test_subscribe_latest_conflates_to_newest(): + bus = FakeAllPubSub() + delivered = [] + first = threading.Event() + gate = threading.Event() + quiet = threading.Event() + + def cb(msg, topic): + delivered.append((topic, msg)) + first.set() + gate.wait(10.0) # block the drain thread while the burst happens + if msg == b"99": + quiet.set() + + unsub = bus.subscribe_latest(cb) + bus.publish("t", b"0") + assert first.wait(5.0) + for i in range(1, 100): + bus.publish("t", str(i).encode()) + gate.set() + + assert quiet.wait(5.0) + # Newest message must arrive; the 98 intermediate ones conflate away. + assert delivered[-1] == ("t", b"99") + assert len(delivered) <= 5 + unsub() + + +def test_subscribe_latest_keeps_latest_per_topic(): + bus = FakeAllPubSub() + delivered = [] + first = threading.Event() + gate = threading.Event() + done = threading.Event() + + def cb(msg, topic): + delivered.append((topic, msg)) + first.set() + gate.wait(10.0) + if {t for t, _ in delivered} >= {"a", "b"} and len(delivered) >= 3: + done.set() + + unsub = bus.subscribe_latest(cb) + bus.publish("a", b"a0") + assert first.wait(5.0) + for i in range(1, 50): + bus.publish("a", f"a{i}".encode()) + bus.publish("b", f"b{i}".encode()) + gate.set() + + assert done.wait(5.0) + latest = dict(delivered) # last write per topic wins + assert latest["a"] == b"a49" + assert latest["b"] == b"b49" + unsub() + + +def test_subscribe_latest_unsubscribe_stops_delivery(): + bus = FakeAllPubSub() + delivered = [] + got_one = threading.Event() + + def cb(msg, topic): + delivered.append(msg) + got_one.set() + + unsub = bus.subscribe_latest(cb) + bus.publish("t", b"1") + assert got_one.wait(5.0) + unsub() + count = len(delivered) + bus.publish("t", b"2") + time.sleep(0.2) + assert len(delivered) == count + + +# SpySources: every message counted, payloads never decoded + + +class _TapCollector: + def __init__(self, topic_str: str, n: int): + self.topic_str = topic_str + self.n = n + self.events: list[tuple[str, int]] = [] + self.done = threading.Event() + + def __call__(self, topic: str, nbytes: int) -> None: + self.events.append((topic, nbytes)) + if len(self.ours()) >= self.n: + self.done.set() + + def ours(self) -> list[tuple[str, int]]: + return [e for e in self.events if e[0] == self.topic_str] + + +def _assert_no_decode(monkeypatch): + """Make any Vector3 decode explode — the spy must never trigger one.""" + + def boom(*a, **k): + raise AssertionError("spy decoded a payload on the hot path") + + monkeypatch.setattr(Vector3, "lcm_decode", staticmethod(boom)) + + +def test_lcm_source_counts_all_without_decoding(monkeypatch): + _assert_no_decode(monkeypatch) + topic = Topic("/spy_e2e", Vector3) + collector = _TapCollector(str(topic), 20) + + src = LCMSpySource() + src.start() + untap = src.tap(collector) + pub = LCMPubSubBase() + pub.start() + try: + time.sleep(0.3) + for _ in range(20): + pub.publish(topic, VEC_BYTES) + time.sleep(0.01) + assert collector.done.wait(10.0) + assert [n for _, n in collector.ours()] == [len(VEC_BYTES)] * 20 + finally: + untap() + src.stop() + pub.stop() + + +def test_lcm_source_counts_undecodable_garbage(monkeypatch): + """Payloads that would crash a decoder must still be counted (proves raw tap).""" + _assert_no_decode(monkeypatch) + topic = Topic("/spy_garbage", Vector3) + collector = _TapCollector(str(topic), 5) + + src = LCMSpySource() + src.start() + untap = src.tap(collector) + pub = LCMPubSubBase() + pub.start() + try: + time.sleep(0.3) + for _ in range(5): + pub.publish(topic, b"\x00garbage-not-lcm") + time.sleep(0.01) + assert collector.done.wait(10.0) + assert [n for _, n in collector.ours()] == [len(b"\x00garbage-not-lcm")] * 5 + finally: + untap() + src.stop() + pub.stop() + + +def test_zenoh_source_counts_all_without_decoding(monkeypatch): + _assert_no_decode(monkeypatch) + topic = Topic("dimos/spy_e2e", Vector3) + collector = _TapCollector(str(topic), 20) + + pool = ZenohSessionPool() + src = ZenohSpySource(session_pool=pool) + src.start() + untap = src.tap(collector) + pub = ZenohPubSubBase(session_pool=pool) + pub.start() + try: + time.sleep(0.5) + for _ in range(20): + pub.publish(topic, VEC_BYTES) + time.sleep(0.01) + assert collector.done.wait(10.0) + assert [n for _, n in collector.ours()] == [len(VEC_BYTES)] * 20 + finally: + untap() + src.stop() + pub.stop() + pool.close_all() + + +# TransportSpy: merging + totals + lifecycle (fake sources, deterministic) + + +class FakeSource: + def __init__(self, name: str): + self.name = name + self.started = False + self.taps: list[Callable[[str, int], None]] = [] + + def start(self) -> None: + self.started = True + + def stop(self) -> None: + self.started = False + + def tap(self, callback): + self.taps.append(callback) + + def untap(): + if callback in self.taps: + self.taps.remove(callback) + + return untap + + def subscribe_decoded(self, topic: str, callback: Callable[[Any], None]): + raise NotImplementedError + + def emit(self, topic: str, nbytes: int) -> None: + for cb in list(self.taps): + cb(topic, nbytes) + + +def test_transport_spy_merges_sources_and_totals(): + a, b = FakeSource("lcm"), FakeSource("zenoh") + spy = TransportSpy(sources=[a, b]) + spy.start() + assert a.started and b.started + + a.emit("/t#geometry_msgs.Twist", 10) + a.emit("/t#geometry_msgs.Twist", 10) + b.emit("dimos/t#geometry_msgs.Twist", 20) + a.emit("/u", 5) + + snap = spy.snapshot() + assert snap[SpyKey("lcm", "/t#geometry_msgs.Twist")].total_msgs == 2 + assert snap[SpyKey("lcm", "/t#geometry_msgs.Twist")].total_bytes == 20 + assert snap[SpyKey("zenoh", "dimos/t#geometry_msgs.Twist")].total_msgs == 1 + assert snap[SpyKey("lcm", "/u")].total_bytes == 5 + assert spy.totals.total_msgs == 4 + assert spy.totals.total_bytes == 45 + + spy.stop() + assert not a.started and not b.started + assert not a.taps and not b.taps # untapped on stop + + +def test_transport_spy_snapshot_is_stable_copy(): + a = FakeSource("lcm") + spy = TransportSpy(sources=[a]) + spy.start() + a.emit("/t", 1) + snap = spy.snapshot() + a.emit("/new_topic_after_snapshot", 1) + assert SpyKey("lcm", "/new_topic_after_snapshot") not in snap # snapshot doesn't mutate + assert SpyKey("lcm", "/new_topic_after_snapshot") in spy.snapshot() + spy.stop() + + +def test_fake_source_satisfies_protocol(): + from dimos.protocol.pubsub.spy import SpySource + + assert isinstance(FakeSource("x"), SpySource) + + +# Lazy decode hook: spec'd now, implemented in a follow-up (stays off hot path) + + +def test_subscribe_decoded_is_not_implemented_in_v1(): + src = LCMSpySource() + with pytest.raises(NotImplementedError): + src.subscribe_decoded("/x#geometry_msgs.Vector3", lambda m: None) diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index d7615a2795..63a0682d05 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -620,13 +620,23 @@ def list_blueprints() -> None: typer.echo(blueprint_name) +@main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) +def spy(ctx: typer.Context) -> None: + """Universal transport spy: topics, rates, sizes across all pubsub transports.""" + from dimos.utils.cli.spy.run_spy import main as spy_main + + sys.argv = ["spy", *ctx.args] + spy_main() + + @main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def lcmspy(ctx: typer.Context) -> None: - """LCM spy tool for monitoring LCM messages.""" - from dimos.utils.cli.lcmspy.run_lcmspy import main as lcmspy_main + """Deprecated alias for `dimos spy --transport lcm`.""" + print("dimos lcmspy is deprecated; use `dimos spy --transport lcm`.", file=sys.stderr) + from dimos.utils.cli.spy.run_spy import main as spy_main - sys.argv = ["lcmspy", *ctx.args] - lcmspy_main() + sys.argv = ["spy", "--transport", "lcm", *ctx.args] + spy_main() @main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) diff --git a/dimos/utils/cli/spy/run_spy.py b/dimos/utils/cli/spy/run_spy.py new file mode 100644 index 0000000000..c33f4c5b5e --- /dev/null +++ b/dimos/utils/cli/spy/run_spy.py @@ -0,0 +1,205 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`dimos spy` TUI: live table of all topics across all pubsub transports. + +Textual app modeled on run_lcmspy.py (DataTable, 0.5s refresh, theme colors, +'q' to quit, optional `spy web` mode via textual-serve). One row per +(transport, topic) from TransportSpy.snapshot(): + + Transport | Topic | Type | Freq (Hz) | Bandwidth | Total | Age + +- Topic/Type come from split_type_suffix(); rows sort by total traffic. +- Age is seconds since TopicStats.last_seen (liveness: dims/greys out stale rows). +- `--transport lcm --transport zenoh` (repeatable) filters sources; default all. +- LCM system config warning: call lcmservice.autoconf(check_only=True) before + entering raw TUI mode, like run_lcmspy does. +""" + +from __future__ import annotations + +import time + +from rich.text import Text +from textual.app import App, ComposeResult +from textual.color import Color +from textual.widgets import DataTable + +from dimos.protocol.pubsub.spy import ( + SpyKey, + TopicStats, + TransportSpy, + default_sources, + split_type_suffix, +) +from dimos.utils.cli import theme +from dimos.utils.human import human_bytes + +# Rows older than this (seconds since last message) render dimmed as "stale". +STALE_AGE = 3.0 +# Window for freq / bandwidth readouts. +STAT_WINDOW = 5.0 + + +def gradient(max_value: float, value: float) -> str: + """Gradient from cyan (low) to yellow (high) using DimOS theme colors.""" + ratio = min(value / max_value, 1.0) if max_value else 0.0 + cyan = Color.parse(theme.CYAN) + yellow = Color.parse(theme.YELLOW) + return cyan.blend(yellow, ratio).hex + + +def topic_text(base: str) -> Text: + """Format a base topic name with DimOS theme colors.""" + if base[:4] == "/rpc": + return Text(base[:4], style=theme.BLUE) + Text(base[4:], style=theme.BRIGHT_WHITE) + return Text(base, style=theme.BRIGHT_WHITE) + + +def _parse_transports(argv: list[str]) -> list[str] | None: + """Parse repeated --transport flags; None means all default sources.""" + transports: list[str] = [] + i = 0 + while i < len(argv): + arg = argv[i] + if arg == "--transport": + i += 1 + if i < len(argv): + transports.append(argv[i]) + elif arg.startswith("--transport="): + transports.append(arg.split("=", 1)[1]) + i += 1 + return transports or None + + +class SpyApp(App): # type: ignore[type-arg] + """A real-time dashboard for all-transport pubsub traffic using Textual.""" + + CSS_PATH = "../dimos.tcss" + + CSS = f""" + Screen {{ + layout: vertical; + background: {theme.BACKGROUND}; + }} + DataTable {{ + height: 2fr; + width: 1fr; + border: solid {theme.BORDER}; + background: {theme.BG}; + scrollbar-size: 0 0; + }} + DataTable > .datatable--header {{ + color: {theme.ACCENT}; + background: transparent; + }} + """ + + refresh_interval: float = 0.5 # seconds + + BINDINGS = [ + ("q", "quit"), + ("ctrl+c", "quit"), + ] + + def __init__(self, transports: list[str] | None = None, **kwargs) -> None: # type: ignore[no-untyped-def] + super().__init__(**kwargs) + sources = default_sources() + if transports is not None: + sources = [s for s in sources if s.name in transports] + # Warn about missing system config before entering TUI raw mode (LCM only). + if any(s.name == "lcm" for s in sources): + from dimos.protocol.service.lcmservice import autoconf + + autoconf(check_only=True) + + self.spy = TransportSpy(sources=sources) + self.spy.start() + self.table: DataTable | None = None # type: ignore[type-arg] + + def compose(self) -> ComposeResult: + self.table = DataTable(zebra_stripes=False, cursor_type=None) # type: ignore[arg-type] + self.table.add_column("Transport") + self.table.add_column("Topic") + self.table.add_column("Type") + self.table.add_column("Freq (Hz)") + self.table.add_column("Bandwidth") + self.table.add_column("Total") + self.table.add_column("Age") + yield self.table + + def on_mount(self) -> None: + self.set_interval(self.refresh_interval, self.refresh_table) + + async def on_unmount(self) -> None: + self.spy.stop() + + def refresh_table(self) -> None: + now = time.time() + snap = self.spy.snapshot() + rows: list[tuple[SpyKey, TopicStats]] = sorted( + snap.items(), key=lambda kv: kv[1].total_bytes, reverse=True + ) + self.table.clear(columns=False) # type: ignore[union-attr] + + for key, stats in rows: + base, msg_type = split_type_suffix(key.topic) + freq = stats.freq(STAT_WINDOW, now) + bps = stats.bytes_per_sec(STAT_WINDOW, now) + age = now - stats.last_seen if stats.last_seen is not None else None + stale = age is not None and age > STALE_AGE + + if stale: + # Liveness: dim the whole row for topics gone quiet. + self.table.add_row( # type: ignore[union-attr] + Text(key.transport, style=theme.DIM), + Text(base, style=theme.DIM), + Text(msg_type or "", style=theme.DIM), + Text(f"{freq:.1f}", style=theme.DIM), + Text(f"{human_bytes(bps)}/s", style=theme.DIM), + Text(human_bytes(stats.total_bytes), style=theme.DIM), + Text(f"{age:.0f}s", style=theme.DIM), + ) + continue + + age_str = f"{age:.1f}s" if age is not None else "-" + self.table.add_row( # type: ignore[union-attr] + Text(key.transport, style=theme.BLUE), + topic_text(base), + Text(msg_type or "", style=theme.BLUE), + Text(f"{freq:.1f}", style=gradient(10, freq)), + Text(f"{human_bytes(bps)}/s", style=gradient(1024 * 3, bps)), + Text(human_bytes(stats.total_bytes)), + Text(age_str, style=theme.ACCENT), + ) + + +def main() -> None: + """Entry point for `dimos spy` (argv: optional 'web', --transport filters).""" + import sys + + argv = sys.argv[1:] + if argv and argv[0] == "web": + import os + + from textual_serve.server import Server # type: ignore[import-not-found] + + server = Server(f"python {os.path.abspath(__file__)}") + server.serve() + else: + SpyApp(transports=_parse_transports(argv)).run() + + +if __name__ == "__main__": + main() diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 960d5e7826..52821b99f3 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -458,7 +458,7 @@ def start(self) -> None: logger.info(f"bridge listening on {pubsub.__class__.__name__}") if hasattr(pubsub, "start"): pubsub.start() - unsub = pubsub.subscribe_all(self._on_message) + unsub = pubsub.subscribe_latest(self._on_message) self.register_disposable(Disposable(unsub)) # Add pubsub stop as disposable diff --git a/dimos/visualization/rerun/test_viewer_integration.py b/dimos/visualization/rerun/test_viewer_integration.py index d6a0ba1e95..a4d5489e35 100644 --- a/dimos/visualization/rerun/test_viewer_integration.py +++ b/dimos/visualization/rerun/test_viewer_integration.py @@ -126,6 +126,9 @@ class ExplicitPubSubOverride: def subscribe_all(self, callback): return lambda: None + def subscribe_latest(self, callback): + return lambda: None + class TestBridgePubsubResolution: def test_legacy_lcm_pubsubs_defers_to_transport_default(self): diff --git a/docs/docs.json b/docs/docs.json index 7cdbc96204..f1cb8c2b86 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -73,7 +73,8 @@ "group": "Transports", "pages": [ "usage/transports/index", - "usage/transports/dds" + "usage/transports/dds", + "usage/transports/spy" ] } ] diff --git a/docs/usage/transports/spy.md b/docs/usage/transports/spy.md new file mode 100644 index 0000000000..d80a05b005 --- /dev/null +++ b/docs/usage/transports/spy.md @@ -0,0 +1,104 @@ +--- +title: "dimos spy" +--- + +`dimos spy` is the universal transport spy: one live view of every topic moving on every +DimOS pubsub transport — names, message rates, bandwidth, sizes, and liveness — regardless +of whether the system runs on LCM, Zenoh, or both. + +> Status: design ratified, implementation in progress. This page is the design doc; +> it becomes the user doc as the tool lands. + +## Motivation + +`dimos lcmspy` only sees LCM. Since the Zenoh transport merged (#2362), a robot can run +with `--transport=zenoh` (default on macOS) and the spy goes blind. Debugging "is data +flowing?" should not require knowing which backend is active — the spy should observe +everything at once. + +## The hard constraint: no decoding on the hot path + +Decoding messages is expensive (LCM decode, image payloads, pointclouds). The spy's job +is *flow* visibility — topic names, Hz, bytes/s, last-seen — none of which needs message +contents. Therefore: + +**The spy never decodes a payload.** Its per-message hot path is +`(topic string, payload length, timestamp)` — an O(1) stats update. + +This falls out of the pubsub layering: encoding/decoding lives in `PubSubEncoderMixin` +subclasses (`LCM`, `Zenoh`), while the base classes (`LCMPubSubBase`, `ZenohPubSubBase`) +move raw bytes. The spy taps the base layer. Message *types* are still displayed for +free because DimOS embeds them in topic strings (`/cmd_vel#geometry_msgs.Twist` on LCM, +`dimos/cmd_vel/geometry_msgs.Twist` as a Zenoh key — both normalize to the `#` form via +`str(Topic)`). + +Per-topic decoded inspection (echo/preview) is a separate, opt-in concern with its own +subscription (`SpySource.subscribe_decoded`), spec'd now and implemented in a follow-up. + +## Architecture + +``` +LCMPubSubBase.subscribe_all(".*") ZenohPubSubBase.subscribe_all("**") + │ raw bytes │ raw bytes + LCMSpySource ("lcm") ZenohSpySource ("zenoh") + │ (topic_str, nbytes) │ (topic_str, nbytes) + └──────────────┬───────────────────────┘ + TransportSpy + dict[SpyKey(transport, topic) → TopicStats] + totals + │ snapshot() + dimos spy TUI + Transport | Topic | Type | Freq | Bandwidth | Total | Age +``` + +- **`SpySource`** (Protocol): one transport's raw firehose. `tap(cb)` delivers + `(topic, nbytes)` for *every* observable message. v1 ships `lcm` and `zenoh` sources; + SHM/ROS/DDS/Redis plug in later by implementing the protocol. +- **`TopicStats`**: sliding-window `(timestamp, nbytes)` history per topic — freq, + bytes/s, avg size over a window, plus lifetime totals and `last_seen` for liveness. + Explicit timestamps keep it deterministic under test. +- **`TransportSpy`**: owns sources, merges taps, serves thread-safe snapshots to the UI. + +Code: `dimos/protocol/pubsub/spy.py` (core), `dimos/utils/cli/spy/run_spy.py` (TUI). + +## The subscribe_all contract fix + +The spy builds on the pubsub layer's subscribe-all abstraction +(`dimos/protocol/pubsub/spec.py`), which exposed a semantic asymmetry: + +- LCM's `subscribe_all` delivered every message (regex `.*`). +- Zenoh's `subscribe_all` **conflated** — latest-wins per topic with a drain thread, + scoped to `dimos/**` — because it was built for its only consumer, the rerun bridge. + +A spy counting rates through a conflating tap silently under-reports. Ratified decision +(Ivan, ticket f6c74d39): **`subscribe_all` is defined non-conflating on every +transport.** Conflation becomes the explicit opt-in `subscribe_latest()` (same +latest-wins + drain-thread semantics, hosted in the spec as a default), and the rerun +bridge migrates to it. Zenoh's `subscribe_all` also widens from `dimos/**` to `**` — +"all topics" means everything the transport can see; the rerun bridge is unaffected +because unknown/untyped keys already fail type resolution and are skipped by its decoder. + +## Decisions (ticket f6c74d39) + +| # | Decision | +|---|----------| +| 1 | `subscribe_all` = every message, non-conflating, on all transports; conflation moves to opt-in `subscribe_latest()`; rerun bridge migrates. Confirmed by Ivan. | +| 2 | v1 sources: LCM + Zenoh, behind the pluggable `SpySource` protocol. SHM/ROS/DDS/Redis later (SHM needs a discovery mechanism first). | +| 3 | `dimos spy` shows **all transports simultaneously**, rows tagged by transport; `dimos lcmspy` becomes a deprecated alias. | +| 4 | Lazy per-topic decode is spec'd (`subscribe_decoded`) but not implemented in v1 — decode stays off the hot path. | +| 5 | This design doc lives in dimos `docs/`. | + +## CLI + +```bash +dimos spy # everything, all transports +dimos spy --transport zenoh # filter to one transport (repeatable flag) +dimos spy web # serve the TUI in a browser (textual-serve) +dimos lcmspy # deprecated alias for: dimos spy --transport lcm +``` + +## Future work + +- SHM source (needs channel discovery), ROS/DDS/Redis sources. +- Per-topic decoded echo/preview from the TUI row (`subscribe_decoded` follow-up). +- Zero-copy zenoh sizing (`len(sample.payload)` without `to_bytes()`), if profiling + shows the copy matters for pointcloud-heavy systems.