fix: Paul's reviews

Author: bogwi

dimos/constants.py

@@ -49,5 +49,10 @@
 # From https://github.com/lcm-proj/lcm.git
 LCM_MAX_CHANNEL_NAME_LENGTH = 63
 
+# First path segment of Zenoh key expressions for DimOS module streams. Used by
+# ModuleCoordinator (dimos + LCM-style "/name") and ZenohPubSub.subscribe_all
+# (wildcard f"{ZENOH_DIMOS_KEY_PREFIX}/**"). Rerun bridge strips this prefix for entity paths.
+ZENOH_DIMOS_KEY_PREFIX = "dimos"
+
 # Default timeout (seconds) for thread.join() during shutdown.
 DEFAULT_THREAD_JOIN_TIMEOUT = 2.0

dimos/core/coordination/module_coordinator.py

@@ -23,6 +23,7 @@
 import threading
 from typing import TYPE_CHECKING, Any, cast
 
+from dimos.constants import ZENOH_DIMOS_KEY_PREFIX
 from dimos.core.coordination.rpyc_server import RpycServer
 from dimos.core.coordination.worker_manager import WorkerManager
 from dimos.core.coordination.worker_manager_docker import WorkerManagerDocker
@@ -557,7 +558,7 @@ def _get_transport_for(blueprint: Blueprint, name: str, stream_type: type) -> Pu
             )
         from dimos.core.transport import ZenohTransport, pZenohTransport
 
-        zenoh_topic = f"dimos{topic}"
+        zenoh_topic = f"{ZENOH_DIMOS_KEY_PREFIX}{topic}"
         transport = (
             pZenohTransport(zenoh_topic)
             if use_pickled
@@ -634,7 +635,9 @@ def _run_configurators(blueprint: Blueprint) -> None:
     from dimos.protocol.service.system_configurator.base import configure_system
     from dimos.protocol.service.system_configurator.lcm_config import lcm_configurators
 
-    lcm_checks = lcm_configurators() if global_config.transport == "lcm" else []
+    # LCM is still used outside merged module transports (CLI side channels, agents, tests).
+    # OS multicast/buffer tuning applies whenever those paths run, not only when transport=lcm.
+    lcm_checks = lcm_configurators()
     configurators = [*lcm_checks, *blueprint.configurator_checks]
 
     try:

dimos/core/test_zenoh_transport.py

@@ -17,24 +17,37 @@
 Tests the conditional logic added to support Zenoh alongside LCM:
 - GlobalConfig transport field
 - _get_transport_for() branching
-- LCM configurator gating
+- LCM system configurators always run with blueprint checks (LCM is used outside transport_map)
+
+Requires the ``zenoh`` extra (``eclipse-zenoh``): ``ZenohTransport`` is only defined
+when that dependency is installed, so this module does not load without it.
 """
 
 from __future__ import annotations
 
+import threading
 from typing import cast
 
+import numpy as np
 from pydantic import ValidationError
 import pytest
 
+from dimos.constants import ZENOH_DIMOS_KEY_PREFIX
 from dimos.core.coordination.blueprints import autoconnect
 from dimos.core.coordination.module_coordinator import _get_transport_for, _run_configurators
 from dimos.core.global_config import GlobalConfig, global_config
 from dimos.core.module import Module
 from dimos.core.stream import In, Out
 from dimos.core.test_utils import retry_until
-from dimos.core.transport import ZENOH_AVAILABLE, LCMTransport, pLCMTransport
+from dimos.core.transport import (
+    ZENOH_AVAILABLE,
+    LCMTransport,
+    ZenohTransport,
+    pLCMTransport,
+    pZenohTransport,
+)
 from dimos.msgs.sensor_msgs.Image import Image
+from dimos.protocol.service.zenohservice import close_all_zenoh_sessions
 
 
 class TypedMsg:
@@ -102,10 +115,7 @@ class TestZenohAvailableGuard:
     def test_zenoh_available_is_bool(self) -> None:
         assert isinstance(ZENOH_AVAILABLE, bool)
 
-    @pytest.mark.skipif(not ZENOH_AVAILABLE, reason="zenoh not installed")
-    def test_zenoh_transport_classes_exist_when_available(self) -> None:
-        from dimos.core.transport import ZenohTransport, pZenohTransport
-
+    def test_zenoh_transport_classes_exist(self) -> None:
         assert ZenohTransport is not None
         assert pZenohTransport is not None
 
@@ -128,33 +138,24 @@ def test_lcm_pickle_transport_returned_for_untyped_when_lcm(self, mocker) -> Non
         transport = _get_transport_for(bp, "untyped_data", UntypedMsg)
         assert isinstance(transport, pLCMTransport)
 
-    @pytest.mark.skipif(not ZENOH_AVAILABLE, reason="zenoh not installed")
     def test_zenoh_transport_returned_when_transport_is_zenoh(self, mocker) -> None:
-        from dimos.core.transport import ZenohTransport
-
         mocker.patch.object(global_config, "transport", "zenoh")
         bp = self._make_blueprint()
         transport = _get_transport_for(bp, "typed_data", TypedMsg)
         assert isinstance(transport, ZenohTransport)
 
-    @pytest.mark.skipif(not ZENOH_AVAILABLE, reason="zenoh not installed")
     def test_zenoh_pickle_transport_returned_for_untyped_when_zenoh(self, mocker) -> None:
-        from dimos.core.transport import pZenohTransport
-
         mocker.patch.object(global_config, "transport", "zenoh")
         bp = self._make_blueprint()
         transport = _get_transport_for(bp, "untyped_data", UntypedMsg)
         assert isinstance(transport, pZenohTransport)
 
-    @pytest.mark.skipif(not ZENOH_AVAILABLE, reason="zenoh not installed")
     def test_zenoh_topic_uses_dimos_prefix(self, mocker) -> None:
-        from dimos.core.transport import pZenohTransport
-
         mocker.patch.object(global_config, "transport", "zenoh")
         bp = self._make_blueprint()
         transport = _get_transport_for(bp, "untyped_data", UntypedMsg)
         assert isinstance(transport, pZenohTransport)
-        assert "dimos/" in transport.topic
+        assert f"{ZENOH_DIMOS_KEY_PREFIX}/" in transport.topic
 
     def test_zenoh_raises_when_not_available(self, mocker) -> None:
         mocker.patch.object(global_config, "transport", "zenoh")
@@ -179,7 +180,7 @@ def test_lcm_configurators_run_when_transport_is_lcm(self, mocker) -> None:
 
         mock_lcm_configs.assert_called_once()
 
-    def test_lcm_configurators_skipped_when_transport_is_zenoh(self, mocker) -> None:
+    def test_lcm_configurators_run_when_transport_is_zenoh(self, mocker) -> None:
         mocker.patch.object(global_config, "transport", "zenoh")
         mock_lcm_configs = mocker.patch(
             "dimos.protocol.service.system_configurator.lcm_config.lcm_configurators",
@@ -190,30 +191,19 @@ def test_lcm_configurators_skipped_when_transport_is_zenoh(self, mocker) -> None
         bp = autoconnect(ProducerModule.blueprint(), ConsumerModule.blueprint())
         _run_configurators(bp)
 
-        mock_lcm_configs.assert_not_called()
+        mock_lcm_configs.assert_called_once()
 
 
-@pytest.mark.skipif(not ZENOH_AVAILABLE, reason="zenoh not installed")
 class TestZenohTransportWrapper:
     """Test ZenohTransport and pZenohTransport broadcast/subscribe lifecycle."""
 
     @pytest.fixture(autouse=True)
     def _clean_sessions(self):
-        from dimos.protocol.service.zenohservice import _sessions
-
         yield
-        for s in _sessions.values():
-            s.close()
-        _sessions.clear()
+        close_all_zenoh_sessions()
 
     def test_zenoh_transport_broadcast_and_subscribe(self) -> None:
-        import threading
-
-        import numpy as np
-
-        from dimos.core.transport import ZenohTransport
-
-        t = ZenohTransport("dimos/test/transport", Image)
+        t = ZenohTransport(f"{ZENOH_DIMOS_KEY_PREFIX}/test/transport", Image)
         t.start()
 
         received = []
@@ -230,11 +220,7 @@ def cb(msg):  # type: ignore[no-untyped-def]
         t.stop()
 
     def test_pzenoh_transport_broadcast_and_subscribe(self) -> None:
-        import threading
-
-        from dimos.core.transport import pZenohTransport
-
-        t = pZenohTransport("dimos/test/pickle_transport")
+        t = pZenohTransport(f"{ZENOH_DIMOS_KEY_PREFIX}/test/pickle_transport")
         t.start()
 
         received = []
@@ -250,18 +236,14 @@ def cb(msg):  # type: ignore[no-untyped-def]
         t.stop()
 
     def test_auto_start_on_broadcast(self) -> None:
-        from dimos.core.transport import pZenohTransport
-
-        t = pZenohTransport("dimos/test/autostart")
+        t = pZenohTransport(f"{ZENOH_DIMOS_KEY_PREFIX}/test/autostart")
         # Don't call start() — broadcast should auto-start
         t.broadcast(None, "test")
         assert t._started
         t.stop()
 
     def test_stop_and_restart(self) -> None:
-        from dimos.core.transport import pZenohTransport
-
-        t = pZenohTransport("dimos/test/restart")
+        t = pZenohTransport(f"{ZENOH_DIMOS_KEY_PREFIX}/test/restart")
         t.start()
         assert t._started
         t.stop()

dimos/protocol/pubsub/impl/test_zenohpubsub.py

@@ -25,25 +25,20 @@
 from dimos.core.test_utils import retry_until
 from dimos.protocol.pubsub.impl.lcmpubsub import Topic
 from dimos.protocol.pubsub.impl.zenohpubsub import ZenohPubSubBase
-from dimos.protocol.service.zenohservice import _sessions
+from dimos.protocol.service.zenohservice import close_all_zenoh_sessions
 
 
 @pytest.fixture()
 def pubsub():
     """Create and start a ZenohPubSubBase instance, clean up after."""
     # Each test gets a fresh session to avoid thread leak detection
-    for session in _sessions.values():
-        session.close()
-    _sessions.clear()
+    close_all_zenoh_sessions()
 
     ps = ZenohPubSubBase()
     ps.start()
     yield ps
     ps.stop()
-    # Close sessions so Zenoh's internal threads are joined
-    for session in _sessions.values():
-        session.close()
-    _sessions.clear()
+    close_all_zenoh_sessions()
 
 
 class TestZenohPubSubBase:

dimos/protocol/pubsub/impl/zenohpubsub.py

@@ -24,6 +24,7 @@
 import threading
 from typing import TYPE_CHECKING, Any
 
+from dimos.constants import ZENOH_DIMOS_KEY_PREFIX
 from dimos.protocol.pubsub.encoders import LCMEncoderMixin, PickleEncoderMixin
 from dimos.protocol.pubsub.impl.lcmpubsub import Topic
 from dimos.protocol.pubsub.spec import AllPubSub
@@ -151,10 +152,10 @@ def on_sample(sample: zenoh.Sample) -> None:
 
         def unsubscribe() -> None:
             nonlocal undeclared
-            if undeclared:
-                return
-            undeclared = True
             with self._subscriber_lock:
+                if undeclared:
+                    return
+                undeclared = True
                 if sub not in self._subscribers:
                     return  # Already removed by stop() — stop() owns the undeclare
                 self._subscribers.remove(sub)
@@ -164,7 +165,7 @@ def unsubscribe() -> None:
 
     def subscribe_all(self, callback: Callable[[bytes, Topic], Any]) -> Callable[[], None]:
         """Subscribe to all dimos key expressions via wildcard."""
-        return self.subscribe(Topic("dimos/**"), callback)
+        return self.subscribe(Topic(f"{ZENOH_DIMOS_KEY_PREFIX}/**"), callback)
 
     def stop(self) -> None:
         """Clean up publishers and subscribers."""

dimos/protocol/service/test_zenohservice.py

@@ -20,17 +20,14 @@
 
 pytest.importorskip("zenoh")
 
-from dimos.protocol.service.zenohservice import ZenohConfig, ZenohService, _sessions
+from dimos.protocol.service.zenohservice import ZenohConfig, ZenohService, close_all_zenoh_sessions
 
 
 @pytest.fixture(autouse=True)
 def _clear_sessions():
     """Clear the global session cache before each test."""
     yield
-    # Close and remove all sessions after each test
-    for session in _sessions.values():
-        session.close()
-    _sessions.clear()
+    close_all_zenoh_sessions()
 
 
 class TestZenohConfig:

dimos/protocol/service/zenohservice.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import atexit
 import json
 import threading
 from typing import Any
@@ -33,6 +34,25 @@
 _sessions_lock = threading.Lock()
 
 
+def close_all_zenoh_sessions() -> None:
+    """Close and clear every cached session in this process.
+
+    Safe to call when no live publishers or subscribers still reference the
+    sessions (call after module ``stop()``). Idempotent if the cache is empty.
+    """
+    with _sessions_lock:
+        to_close = list(_sessions.values())
+        _sessions.clear()
+    for session in to_close:
+        try:
+            session.close()  # type: ignore[no-untyped-call]
+        except Exception:
+            logger.error("Error closing Zenoh session", exc_info=True)
+
+
+atexit.register(close_all_zenoh_sessions)
+
+
 class ZenohConfig(BaseConfig):
     """Configuration for Zenoh service."""
 
@@ -83,4 +103,5 @@ def session(self) -> zenoh.Session:
 __all__ = [
     "ZenohConfig",
     "ZenohService",
+    "close_all_zenoh_sessions",
 ]

dimos/visualization/rerun/bridge.py

@@ -39,6 +39,7 @@
 from toolz import pipe  # type: ignore[import-untyped]
 import typer
 
+from dimos.constants import ZENOH_DIMOS_KEY_PREFIX
 from dimos.core.core import rpc
 from dimos.core.global_config import global_config
 from dimos.core.module import Module, ModuleConfig
@@ -321,9 +322,10 @@ def _get_entity_path(self, topic: Any) -> str:
         if isinstance(raw, str):
             topic_str = raw
         topic_str = topic_str.split("#")[0]
-        # Strip Zenoh key prefix (dimos/) to match LCM entity paths
-        if topic_str.startswith("dimos/"):
-            topic_str = "/" + topic_str.removeprefix("dimos/")
+        # Strip Zenoh key root to match LCM-style entity paths.
+        _prefix = f"{ZENOH_DIMOS_KEY_PREFIX}/"
+        if topic_str.startswith(_prefix):
+            topic_str = "/" + topic_str.removeprefix(_prefix)
         return f"{self.config.entity_prefix}{topic_str}"
 
     def _on_message(self, msg: Any, topic: Any) -> None: