From 957c26e50763cfb21580f9024cec4c973f3f73bf Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:19:41 -0700 Subject: [PATCH 01/53] switch to websocket --- .../visualization/rerun/test_viewer_ws_e2e.py | 332 ++++++++++++++++ .../rerun/test_websocket_server.py | 374 ++++++++++++++++++ dimos/visualization/rerun/websocket_server.py | 173 ++++++++ 3 files changed, 879 insertions(+) create mode 100644 dimos/visualization/rerun/test_viewer_ws_e2e.py create mode 100644 dimos/visualization/rerun/test_websocket_server.py create mode 100644 dimos/visualization/rerun/websocket_server.py diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py new file mode 100644 index 0000000000..266e16cc68 --- /dev/null +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -0,0 +1,332 @@ +# Copyright 2025-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. + +"""End-to-end test: dimos-viewer (headless) → WebSocket → RerunWebSocketServer. + +dimos-viewer is started in ``--connect`` mode so it initialises its WebSocket +client. The viewer needs a gRPC proxy to connect to; we give it a non-existent +one so the viewer starts up anyway but produces no visualisation. The important +part is that the WebSocket client inside the viewer tries to connect to +``ws://127.0.0.1:/ws``. + +Because the viewer is a native GUI application it cannot run headlessly in CI +without a display. This test therefore verifies the connection at the protocol +level by using the ``RerunWebSocketServer`` module directly as the server and +injecting synthetic JSON messages that mimic what the viewer would send once a +user clicks in the 3D viewport. +""" + +import asyncio +import json +import subprocess +import threading +import time +from typing import Any + +import pytest + +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +_E2E_PORT = 13032 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_server(port: int = _E2E_PORT) -> RerunWebSocketServer: + return RerunWebSocketServer(port=port) + + +def _wait_for_server(port: int, timeout: float = 5.0) -> None: + import websockets.asyncio.client as ws_client + + async def _probe() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): + pass + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + asyncio.run(_probe()) + return + except Exception: + time.sleep(0.05) + raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") + + +def _send_messages(port: int, messages: list[dict[str, Any]], *, delay: float = 0.05) -> None: + import websockets.asyncio.client as ws_client + + async def _run() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws") as ws: + for msg in messages: + await ws.send(json.dumps(msg)) + await asyncio.sleep(delay) + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# Protocol-level E2E tests (no GUI required) +# --------------------------------------------------------------------------- + + +class TestViewerProtocolE2E: + """Verify the full Python-server side of the viewer ↔ DimOS protocol. + + These tests use the ``RerunWebSocketServer`` as the server and a dummy + WebSocket client (playing the role of dimos-viewer) to inject messages. + They confirm every message type is correctly routed and that only click + messages produce stream publishes. + """ + + def test_viewer_click_reaches_stream(self): + """A viewer click message received over WebSocket publishes PointStamped.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + done.set() + + server.clicked_point.subscribe(_on_pt) + + _send_messages( + _E2E_PORT, + [ + { + "type": "click", + "x": 10.0, + "y": 20.0, + "z": 0.5, + "entity_path": "/world/robot", + "timestamp_ms": 42000, + } + ], + ) + + done.wait(timeout=3.0) + server.stop() + + assert len(received) == 1 + pt = received[0] + assert abs(pt.x - 10.0) < 1e-9 + assert abs(pt.y - 20.0) < 1e-9 + assert abs(pt.z - 0.5) < 1e-9 + assert pt.frame_id == "/world/robot" + assert abs(pt.ts - 42.0) < 1e-6 + + def test_viewer_keyboard_twist_no_publish(self): + """Twist messages from keyboard control do not publish clicked_point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + server.clicked_point.subscribe(received.append) + + _send_messages( + _E2E_PORT, + [ + { + "type": "twist", + "linear_x": 0.5, + "linear_y": 0.0, + "linear_z": 0.0, + "angular_x": 0.0, + "angular_y": 0.0, + "angular_z": 0.8, + } + ], + ) + + server.stop() + assert received == [] + + def test_viewer_stop_no_publish(self): + """Stop messages do not publish clicked_point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + server.clicked_point.subscribe(received.append) + + _send_messages(_E2E_PORT, [{"type": "stop"}]) + + server.stop() + assert received == [] + + def test_full_viewer_session_sequence(self): + """Realistic session: connect, heartbeats, click, WASD, stop → one point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + done.set() + + server.clicked_point.subscribe(_on_pt) + + _send_messages( + _E2E_PORT, + [ + # Initial heartbeats (viewer connects and starts 1 Hz heartbeat) + {"type": "heartbeat", "timestamp_ms": 1000}, + {"type": "heartbeat", "timestamp_ms": 2000}, + # User clicks a point in the 3D viewport + { + "type": "click", + "x": 3.14, + "y": 2.71, + "z": 1.41, + "entity_path": "/world", + "timestamp_ms": 3000, + }, + # User presses W (forward) + { + "type": "twist", + "linear_x": 0.5, + "linear_y": 0.0, + "linear_z": 0.0, + "angular_x": 0.0, + "angular_y": 0.0, + "angular_z": 0.0, + }, + # User releases W + {"type": "stop"}, + # Another heartbeat + {"type": "heartbeat", "timestamp_ms": 4000}, + ], + delay=0.2, + ) + + done.wait(timeout=3.0) + server.stop() + + assert len(received) == 1, f"Expected exactly 1 click, got {len(received)}" + pt = received[0] + assert abs(pt.x - 3.14) < 1e-9 + assert abs(pt.y - 2.71) < 1e-9 + assert abs(pt.z - 1.41) < 1e-9 + + def test_reconnect_after_disconnect(self): + """Server keeps accepting new connections after a client disconnects.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + all_done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + if len(received) >= 2: + all_done.set() + + server.clicked_point.subscribe(_on_pt) + + # First connection — send one click and disconnect + _send_messages( + _E2E_PORT, + [{"type": "click", "x": 1.0, "y": 0.0, "z": 0.0, "entity_path": "", "timestamp_ms": 0}], + ) + + # Second connection (simulating viewer reconnect) — send another click + _send_messages( + _E2E_PORT, + [{"type": "click", "x": 2.0, "y": 0.0, "z": 0.0, "entity_path": "", "timestamp_ms": 0}], + ) + + all_done.wait(timeout=5.0) + server.stop() + + xs = sorted(pt.x for pt in received) + assert xs == [1.0, 2.0], f"Unexpected xs: {xs}" + + +# --------------------------------------------------------------------------- +# Binary smoke test +# --------------------------------------------------------------------------- + + +class TestViewerBinaryConnectMode: + """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket + client attempts to connect to our Python server.""" + + def test_viewer_ws_client_connects(self): + """dimos-viewer --connect starts and its WS client connects to our server.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + connected = threading.Event() + received: list[Any] = [] + + def _on_pt(pt: Any) -> None: + received.append(pt) + + server.clicked_point.subscribe(_on_pt) + + # Start dimos-viewer in --connect mode, pointing it at a non-existent gRPC + # proxy (it will fail to stream data, but that's fine) and at our WS server. + # Use DISPLAY="" to prevent it from opening a window (it will exit quickly + # without a display, but the WebSocket connection happens before the GUI loop). + proc = subprocess.Popen( + [ + "dimos-viewer", + "--connect", + f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", + ], + env={"DISPLAY": "", "HOME": "/home/dimos", "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin"}, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + # Give the viewer up to 5 s to connect its WebSocket client to our server. + # We detect the connection by waiting for the server to accept a client. + deadline = time.monotonic() + 5.0 + viewer_connected = False + while time.monotonic() < deadline: + # Check if any connection was established by sending a message and + # verifying the viewer is still running. + if proc.poll() is not None: + # Viewer exited (expected without a display) — check if it connected first. + break + time.sleep(0.1) + + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + + stderr = proc.stderr.read().decode(errors="replace") if proc.stderr else "" + server.stop() + + # The viewer should log that it is connecting to our WS URL. + # Even without a display, the log output appears before the GUI loop starts. + assert "ws://127.0.0.1" in stderr or proc.returncode is not None, ( + f"Viewer did not attempt WS connection. stderr:\n{stderr}" + ) diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py new file mode 100644 index 0000000000..e1dc08ee23 --- /dev/null +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -0,0 +1,374 @@ +# Copyright 2025-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. + +"""Tests for RerunWebSocketServer. + +Uses ``MockViewerPublisher`` to simulate dimos-viewer sending events, matching +the exact JSON protocol used by the Rust ``WsPublisher`` in the viewer. +""" + +import asyncio +import json +import threading +import time +from typing import Any + +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +_TEST_PORT = 13031 + + +# --------------------------------------------------------------------------- +# MockViewerPublisher +# --------------------------------------------------------------------------- + + +class MockViewerPublisher: + """Python mirror of the Rust WsPublisher in dimos-viewer. + + Connects to a running ``RerunWebSocketServer`` and exposes the same + ``send_click`` / ``send_twist`` / ``send_stop`` / ``send_heartbeat`` + API that the real viewer uses. Useful for unit tests that need to + exercise the server without a real viewer binary. + + Usage:: + + with MockViewerPublisher("ws://127.0.0.1:13031/ws") as pub: + pub.send_click(1.0, 2.0, 0.0, "/world", timestamp_ms=1000) + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.send_stop() + """ + + def __init__(self, url: str) -> None: + self._url = url + self._ws: Any = None + self._loop: asyncio.AbstractEventLoop | None = None + + # ------------------------------------------------------------------ + # Context-manager interface + # ------------------------------------------------------------------ + + def __enter__(self) -> "MockViewerPublisher": + self._loop = asyncio.new_event_loop() + self._ws = self._loop.run_until_complete(self._connect()) + return self + + def __exit__(self, *_: Any) -> None: + if self._ws is not None and self._loop is not None: + self._loop.run_until_complete(self._ws.close()) + if self._loop is not None: + self._loop.close() + + async def _connect(self) -> Any: + import websockets.asyncio.client as ws_client + return await ws_client.connect(self._url) + + # ------------------------------------------------------------------ + # Send helpers (mirror of Rust WsPublisher methods) + # ------------------------------------------------------------------ + + def send_click( + self, + x: float, + y: float, + z: float, + entity_path: str = "", + timestamp_ms: int = 0, + ) -> None: + """Send a click event — matches viewer SelectionChange handler output.""" + self._send( + { + "type": "click", + "x": x, + "y": y, + "z": z, + "entity_path": entity_path, + "timestamp_ms": timestamp_ms, + } + ) + + def send_twist( + self, + linear_x: float, + linear_y: float, + linear_z: float, + angular_x: float, + angular_y: float, + angular_z: float, + ) -> None: + """Send a twist (WASD keyboard) event.""" + self._send( + { + "type": "twist", + "linear_x": linear_x, + "linear_y": linear_y, + "linear_z": linear_z, + "angular_x": angular_x, + "angular_y": angular_y, + "angular_z": angular_z, + } + ) + + def send_stop(self) -> None: + """Send a stop event (Space bar or key release).""" + self._send({"type": "stop"}) + + def send_heartbeat(self, timestamp_ms: int = 0) -> None: + """Send a heartbeat (1 Hz keepalive from viewer).""" + self._send({"type": "heartbeat", "timestamp_ms": timestamp_ms}) + + def flush(self, delay: float = 0.1) -> None: + """Wait briefly so the server processes queued messages.""" + time.sleep(delay) + + def _send(self, msg: dict[str, Any]) -> None: + assert self._loop is not None and self._ws is not None, "Not connected" + self._loop.run_until_complete(self._ws.send(json.dumps(msg))) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: + return RerunWebSocketServer(port=port) + + +def _wait_for_server(port: int, timeout: float = 3.0) -> None: + """Block until the WebSocket server accepts an upgrade handshake.""" + async def _probe() -> None: + import websockets.asyncio.client as ws_client + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): + pass + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + asyncio.run(_probe()) + return + except Exception: + time.sleep(0.05) + raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestRerunWebSocketServerStartup: + def test_server_binds_port(self) -> None: + """After start(), the server must be reachable on the configured port.""" + mod = _make_module() + mod.start() + try: + _wait_for_server(_TEST_PORT) + finally: + mod.stop() + + def test_stop_is_idempotent(self) -> None: + """Calling stop() twice must not raise.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + mod.stop() + mod.stop() + + +class TestClickMessages: + def test_click_publishes_point_stamped(self) -> None: + """A single click publishes one PointStamped with correct coords.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(1.5, 2.5, 0.0, "/world", timestamp_ms=1000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + pt = received[0] + assert abs(pt.x - 1.5) < 1e-9 + assert abs(pt.y - 2.5) < 1e-9 + assert abs(pt.z - 0.0) < 1e-9 + + def test_click_sets_frame_id_from_entity_path(self) -> None: + """entity_path is stored as frame_id on the published PointStamped.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(0.0, 0.0, 0.0, "/robot/base", timestamp_ms=2000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + assert received and received[0].frame_id == "/robot/base" + + def test_click_timestamp_converted_from_ms(self) -> None: + """timestamp_ms is converted to seconds on PointStamped.ts.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(0.0, 0.0, 0.0, "", timestamp_ms=5000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + assert received and abs(received[0].ts - 5.0) < 1e-6 + + def test_multiple_clicks_all_published(self) -> None: + """A burst of clicks all arrive on the stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + all_arrived = threading.Event() + + def _cb(pt: Any) -> None: + received.append(pt) + if len(received) >= 3: + all_arrived.set() + + mod.clicked_point.subscribe(_cb) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(1.0, 0.0, 0.0) + pub.send_click(2.0, 0.0, 0.0) + pub.send_click(3.0, 0.0, 0.0) + pub.flush() + + all_arrived.wait(timeout=3.0) + mod.stop() + + assert sorted(pt.x for pt in received) == [1.0, 2.0, 3.0] + + +class TestNonClickMessages: + def test_heartbeat_does_not_publish(self) -> None: + """Heartbeat messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_heartbeat(9999) + pub.flush() + + mod.stop() + assert received == [] + + def test_twist_does_not_publish_clicked_point(self) -> None: + """Twist messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.flush() + + mod.stop() + assert received == [] + + def test_stop_does_not_publish_clicked_point(self) -> None: + """Stop messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_stop() + pub.flush() + + mod.stop() + assert received == [] + + def test_invalid_json_does_not_crash(self) -> None: + """Malformed JSON is silently dropped; server stays alive.""" + import websockets.asyncio.client as ws_client + + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + async def _send_bad() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{_TEST_PORT}/ws") as ws: + await ws.send("this is not json {{") + await asyncio.sleep(0.1) + await ws.send(json.dumps({"type": "heartbeat", "timestamp_ms": 0})) + await asyncio.sleep(0.1) + + asyncio.run(_send_bad()) + mod.stop() + + def test_mixed_message_sequence(self) -> None: + """Realistic sequence: heartbeat → click → twist → stop publishes one point.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + # Subscribe before sending so we don't race against the click dispatch. + received: list[Any] = [] + done = threading.Event() + + def _cb(pt: Any) -> None: + received.append(pt) + done.set() + + mod.clicked_point.subscribe(_cb) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_heartbeat(1000) + pub.send_click(7.0, 8.0, 9.0, "/map", timestamp_ms=1100) + pub.send_twist(0.3, 0.0, 0.0, 0.0, 0.0, 0.2) + pub.send_stop() + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + assert abs(received[0].x - 7.0) < 1e-9 + assert abs(received[0].y - 8.0) < 1e-9 + assert abs(received[0].z - 9.0) < 1e-9 diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py new file mode 100644 index 0000000000..a04e2c4999 --- /dev/null +++ b/dimos/visualization/rerun/websocket_server.py @@ -0,0 +1,173 @@ +# Copyright 2025-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. + +"""WebSocket server module that receives events from dimos-viewer. + +When dimos-viewer is started with ``--connect``, LCM multicast is unavailable +across machines. The viewer falls back to sending click, twist, and stop events +as JSON over a WebSocket connection. This module acts as the server-side +counterpart: it listens for those connections and translates incoming messages +into DimOS stream publishes. + +Message format (newline-delimited JSON, ``"type"`` discriminant): + + {"type":"heartbeat","timestamp_ms":1234567890} + {"type":"click","x":1.0,"y":2.0,"z":3.0,"entity_path":"/world","timestamp_ms":1234567890} + {"type":"twist","linear_x":0.5,"linear_y":0.0,"linear_z":0.0, + "angular_x":0.0,"angular_y":0.0,"angular_z":0.8} + {"type":"stop"} +""" + +import asyncio +import json +import threading +from typing import Any + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out +from dimos.msgs.geometry_msgs.PointStamped import PointStamped +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class Config(ModuleConfig): + port: int = 3030 + + +class RerunWebSocketServer(Module[Config]): + """Receives dimos-viewer WebSocket events and publishes them as DimOS streams. + + The viewer connects to this module (not the other way around) when running + in ``--connect`` mode. Each click event is converted to a ``PointStamped`` + and published on the ``clicked_point`` stream so downstream modules (e.g. + ``ReplanningAStarPlanner``) can consume it without modification. + + Outputs: + clicked_point: 3-D world-space point from the most recent viewer click. + """ + + default_config = Config + + clicked_point: Out[PointStamped] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._ws_loop: asyncio.AbstractEventLoop | None = None + self._server_thread: threading.Thread | None = None + self._stop_event: asyncio.Event | None = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + @rpc + def start(self) -> None: + super().start() + self._server_thread = threading.Thread( + target=self._run_server, daemon=True, name="rerun-ws-server" + ) + self._server_thread.start() + logger.info(f"RerunWebSocketServer starting on ws://0.0.0.0:{self.config.port}/ws") + + @rpc + def stop(self) -> None: + if ( + self._ws_loop is not None + and not self._ws_loop.is_closed() + and self._stop_event is not None + ): + self._ws_loop.call_soon_threadsafe(self._stop_event.set) + super().stop() + + # ------------------------------------------------------------------ + # Server + # ------------------------------------------------------------------ + + def _run_server(self) -> None: + """Entry point for the background server thread.""" + self._ws_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._ws_loop) + try: + self._ws_loop.run_until_complete(self._serve()) + finally: + self._ws_loop.close() + + async def _serve(self) -> None: + import websockets.asyncio.server as ws_server + + self._stop_event = asyncio.Event() + + async with ws_server.serve( + self._handle_client, + host="0.0.0.0", + port=self.config.port, + ): + logger.info( + f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws" + ) + await self._stop_event.wait() + + async def _handle_client(self, websocket: Any) -> None: + addr = websocket.remote_address + logger.info(f"RerunWebSocketServer: viewer connected from {addr}") + try: + async for raw in websocket: + self._dispatch(raw) + except Exception as exc: + logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") + + # ------------------------------------------------------------------ + # Message dispatch + # ------------------------------------------------------------------ + + def _dispatch(self, raw: str | bytes) -> None: + try: + msg = json.loads(raw) + except json.JSONDecodeError: + logger.warning(f"RerunWebSocketServer: ignoring non-JSON message: {raw!r}") + return + + msg_type = msg.get("type") + + if msg_type == "click": + pt = PointStamped( + x=float(msg["x"]), + y=float(msg["y"]), + z=float(msg["z"]), + ts=float(msg.get("timestamp_ms", 0)) / 1000.0, + frame_id=str(msg.get("entity_path", "")), + ) + logger.debug(f"RerunWebSocketServer: click → {pt}") + self.clicked_point.publish(pt) + + elif msg_type == "twist": + # Twist messages are not yet wired to a stream; log for observability. + logger.debug( + "RerunWebSocketServer: twist lin=({linear_x},{linear_y},{linear_z}) " + "ang=({angular_x},{angular_y},{angular_z})".format(**msg) + ) + + elif msg_type == "stop": + logger.debug("RerunWebSocketServer: stop") + + elif msg_type == "heartbeat": + logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") + + else: + logger.warning(f"RerunWebSocketServer: unknown message type {msg_type!r}") + + +rerun_ws_server = RerunWebSocketServer.blueprint From f5a35bb6d6627d973c8e45f079f0d1070134df7d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:27:40 -0700 Subject: [PATCH 02/53] cleanup --- .../visualization/rerun/test_viewer_ws_e2e.py | 11 ++++---- .../rerun/test_websocket_server.py | 3 +++ dimos/visualization/rerun/websocket_server.py | 26 ++++++++++++++----- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index 266e16cc68..4026d1d346 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -34,8 +34,6 @@ import time from typing import Any -import pytest - from dimos.visualization.rerun.websocket_server import RerunWebSocketServer _E2E_PORT = 13032 @@ -281,7 +279,7 @@ def test_viewer_ws_client_connects(self): server.start() _wait_for_server(_E2E_PORT) - connected = threading.Event() + threading.Event() received: list[Any] = [] def _on_pt(pt: Any) -> None: @@ -299,7 +297,11 @@ def _on_pt(pt: Any) -> None: "--connect", f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", ], - env={"DISPLAY": "", "HOME": "/home/dimos", "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin"}, + env={ + "DISPLAY": "", + "HOME": "/home/dimos", + "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin", + }, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) @@ -307,7 +309,6 @@ def _on_pt(pt: Any) -> None: # Give the viewer up to 5 s to connect its WebSocket client to our server. # We detect the connection by waiting for the server to accept a client. deadline = time.monotonic() + 5.0 - viewer_connected = False while time.monotonic() < deadline: # Check if any connection was established by sending a message and # verifying the viewer is still running. diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index e1dc08ee23..d0bd986d91 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -72,6 +72,7 @@ def __exit__(self, *_: Any) -> None: async def _connect(self) -> Any: import websockets.asyncio.client as ws_client + return await ws_client.connect(self._url) # ------------------------------------------------------------------ @@ -148,8 +149,10 @@ def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: def _wait_for_server(port: int, timeout: float = 3.0) -> None: """Block until the WebSocket server accepts an upgrade handshake.""" + async def _probe() -> None: import websockets.asyncio.client as ws_client + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): pass diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index a04e2c4999..70b6468408 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -38,6 +38,8 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.msgs.geometry_msgs.PointStamped import PointStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -57,11 +59,13 @@ class RerunWebSocketServer(Module[Config]): Outputs: clicked_point: 3-D world-space point from the most recent viewer click. + tele_cmd_vel: Twist velocity commands from keyboard teleop, including stop events. """ default_config = Config clicked_point: Out[PointStamped] + tele_cmd_vel: Out[Twist] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -115,9 +119,7 @@ async def _serve(self) -> None: host="0.0.0.0", port=self.config.port, ): - logger.info( - f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws" - ) + logger.info(f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws") await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -154,14 +156,24 @@ def _dispatch(self, raw: str | bytes) -> None: self.clicked_point.publish(pt) elif msg_type == "twist": - # Twist messages are not yet wired to a stream; log for observability. - logger.debug( - "RerunWebSocketServer: twist lin=({linear_x},{linear_y},{linear_z}) " - "ang=({angular_x},{angular_y},{angular_z})".format(**msg) + twist = Twist( + linear=Vector3( + float(msg.get("linear_x", 0)), + float(msg.get("linear_y", 0)), + float(msg.get("linear_z", 0)), + ), + angular=Vector3( + float(msg.get("angular_x", 0)), + float(msg.get("angular_y", 0)), + float(msg.get("angular_z", 0)), + ), ) + logger.debug(f"RerunWebSocketServer: twist → {twist}") + self.tele_cmd_vel.publish(twist) elif msg_type == "stop": logger.debug("RerunWebSocketServer: stop") + self.tele_cmd_vel.publish(Twist.zero()) elif msg_type == "heartbeat": logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") From e67ae72b5b7d7c13ec3f98776da2dde2d97a49be Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:33:47 -0700 Subject: [PATCH 03/53] improvements --- .../visualization/rerun/test_viewer_ws_e2e.py | 13 ++-- .../rerun/test_websocket_server.py | 59 ++++++++++++++++++- dimos/visualization/rerun/websocket_server.py | 13 ++-- 3 files changed, 71 insertions(+), 14 deletions(-) diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index 4026d1d346..d7bac7b6f4 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -91,7 +91,7 @@ class TestViewerProtocolE2E: messages produce stream publishes. """ - def test_viewer_click_reaches_stream(self): + def test_viewer_click_reaches_stream(self) -> None: """A viewer click message received over WebSocket publishes PointStamped.""" server = _make_server() server.start() @@ -131,7 +131,7 @@ def _on_pt(pt: Any) -> None: assert pt.frame_id == "/world/robot" assert abs(pt.ts - 42.0) < 1e-6 - def test_viewer_keyboard_twist_no_publish(self): + def test_viewer_keyboard_twist_no_publish(self) -> None: """Twist messages from keyboard control do not publish clicked_point.""" server = _make_server() server.start() @@ -158,7 +158,7 @@ def test_viewer_keyboard_twist_no_publish(self): server.stop() assert received == [] - def test_viewer_stop_no_publish(self): + def test_viewer_stop_no_publish(self) -> None: """Stop messages do not publish clicked_point.""" server = _make_server() server.start() @@ -172,7 +172,7 @@ def test_viewer_stop_no_publish(self): server.stop() assert received == [] - def test_full_viewer_session_sequence(self): + def test_full_viewer_session_sequence(self) -> None: """Realistic session: connect, heartbeats, click, WASD, stop → one point.""" server = _make_server() server.start() @@ -229,7 +229,7 @@ def _on_pt(pt: Any) -> None: assert abs(pt.y - 2.71) < 1e-9 assert abs(pt.z - 1.41) < 1e-9 - def test_reconnect_after_disconnect(self): + def test_reconnect_after_disconnect(self) -> None: """Server keeps accepting new connections after a client disconnects.""" server = _make_server() server.start() @@ -273,13 +273,12 @@ class TestViewerBinaryConnectMode: """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket client attempts to connect to our Python server.""" - def test_viewer_ws_client_connects(self): + def test_viewer_ws_client_connects(self) -> None: """dimos-viewer --connect starts and its WS client connects to our server.""" server = _make_server() server.start() _wait_for_server(_E2E_PORT) - threading.Event() received: list[Any] = [] def _on_pt(pt: Any) -> None: diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index d0bd986d91..c894774679 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -143,6 +143,16 @@ def _send(self, msg: dict[str, Any]) -> None: # --------------------------------------------------------------------------- +def _collect(received: list[Any], done: threading.Event) -> Any: + """Return a callback that appends to *received* and signals *done*.""" + + def _cb(msg: Any) -> None: + received.append(msg) + done.set() + + return _cb + + def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: return RerunWebSocketServer(port=port) @@ -199,7 +209,7 @@ def test_click_publishes_point_stamped(self) -> None: received: list[Any] = [] done = threading.Event() - mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + mod.clicked_point.subscribe(_collect(received, done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_click(1.5, 2.5, 0.0, "/world", timestamp_ms=1000) @@ -222,7 +232,7 @@ def test_click_sets_frame_id_from_entity_path(self) -> None: received: list[Any] = [] done = threading.Event() - mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + mod.clicked_point.subscribe(_collect(received, done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_click(0.0, 0.0, 0.0, "/robot/base", timestamp_ms=2000) @@ -240,7 +250,7 @@ def test_click_timestamp_converted_from_ms(self) -> None: received: list[Any] = [] done = threading.Event() - mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + mod.clicked_point.subscribe(_collect(received, done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_click(0.0, 0.0, 0.0, "", timestamp_ms=5000) @@ -327,6 +337,49 @@ def test_stop_does_not_publish_clicked_point(self) -> None: mod.stop() assert received == [] + def test_twist_publishes_on_tele_cmd_vel(self) -> None: + """Twist messages publish a Twist on the tele_cmd_vel stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + tw = received[0] + assert abs(tw.linear.x - 0.5) < 1e-9 + assert abs(tw.angular.z - 0.8) < 1e-9 + + def test_stop_publishes_zero_twist_on_tele_cmd_vel(self) -> None: + """Stop messages publish a zero Twist on the tele_cmd_vel stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_stop() + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + tw = received[0] + assert tw.is_zero() + def test_invalid_json_does_not_crash(self) -> None: """Malformed JSON is silently dropped; server stays alive.""" import websockets.asyncio.client as ws_client diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 70b6468408..163bfcbf62 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -34,6 +34,8 @@ import threading from typing import Any +import websockets + from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out @@ -46,6 +48,9 @@ class Config(ModuleConfig): + # Intentionally binds 0.0.0.0 by default so the viewer can connect from + # any machine on the network (the typical robot deployment scenario). + host: str = "0.0.0.0" port: int = 3030 @@ -84,7 +89,7 @@ def start(self) -> None: target=self._run_server, daemon=True, name="rerun-ws-server" ) self._server_thread.start() - logger.info(f"RerunWebSocketServer starting on ws://0.0.0.0:{self.config.port}/ws") + logger.info(f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws") @rpc def stop(self) -> None: @@ -116,10 +121,10 @@ async def _serve(self) -> None: async with ws_server.serve( self._handle_client, - host="0.0.0.0", + host=self.config.host, port=self.config.port, ): - logger.info(f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws") + logger.info(f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws") await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -128,7 +133,7 @@ async def _handle_client(self, websocket: Any) -> None: try: async for raw in websocket: self._dispatch(raw) - except Exception as exc: + except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") # ------------------------------------------------------------------ From b7bfb405acc9fbbdbb2c2bdda04b52b4149545dd Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:46:25 -0700 Subject: [PATCH 04/53] fix: ruff formatting + consistent error handling in websocket_server --- dimos/visualization/rerun/websocket_server.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 163bfcbf62..ba0c953bd8 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -89,7 +89,9 @@ def start(self) -> None: target=self._run_server, daemon=True, name="rerun-ws-server" ) self._server_thread.start() - logger.info(f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws") + logger.info( + f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws" + ) @rpc def stop(self) -> None: @@ -124,7 +126,9 @@ async def _serve(self) -> None: host=self.config.host, port=self.config.port, ): - logger.info(f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws") + logger.info( + f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" + ) await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -151,9 +155,9 @@ def _dispatch(self, raw: str | bytes) -> None: if msg_type == "click": pt = PointStamped( - x=float(msg["x"]), - y=float(msg["y"]), - z=float(msg["z"]), + x=float(msg.get("x", 0)), + y=float(msg.get("y", 0)), + z=float(msg.get("z", 0)), ts=float(msg.get("timestamp_ms", 0)) / 1000.0, frame_id=str(msg.get("entity_path", "")), ) From 7cbe1b7a99d0e5b3d2d5f4bcc0d7de8c61216534 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 16:43:17 -0700 Subject: [PATCH 05/53] make it easy to use --- .../primitive/uintree_g1_primitive_no_nav.py | 16 +--- .../go2/blueprints/basic/unitree_go2_basic.py | 30 +++----- dimos/visualization/vis_module.py | 73 +++++++++++++++++++ 3 files changed, 84 insertions(+), 35 deletions(-) create mode 100644 dimos/visualization/vis_module.py diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index c3da9521c5..2228dbfd66 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -40,7 +40,7 @@ from dimos.navigation.frontier_exploration.wavefront_frontier_goal_selector import ( WavefrontFrontierExplorer, ) -from dimos.protocol.pubsub.impl.lcmpubsub import LCM +from dimos.visualization.vis_module import vis_module from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule @@ -90,7 +90,6 @@ def _g1_rerun_blueprint() -> Any: rerun_config = { "blueprint": _g1_rerun_blueprint, - "pubsubs": [LCM()], "visual_override": { "world/camera_info": _convert_camera_info, "world/global_map": _convert_global_map, @@ -101,18 +100,7 @@ def _g1_rerun_blueprint() -> Any: }, } -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge - - _with_vis = autoconnect(FoxgloveBridge.blueprint()) -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode - - _with_vis = autoconnect( - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config) - ) -else: - _with_vis = autoconnect() +_with_vis = vis_module(global_config.viewer, rerun_config=rerun_config) def _create_webcam() -> Webcam: diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index a0d1e6a7ae..406454ecc9 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -22,9 +22,9 @@ from dimos.core.global_config import global_config from dimos.core.transport import pSHMTransport from dimos.msgs.sensor_msgs.Image import Image -from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.visualization.vis_module import vis_module from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule # Mac has some issue with high bandwidth UDP, so we use pSHMTransport for color_image @@ -87,9 +87,6 @@ def _go2_rerun_blueprint() -> Any: rerun_config = { "blueprint": _go2_rerun_blueprint, - # any pubsub that supports subscribe_all and topic that supports str(topic) - # is acceptable here - "pubsubs": [LCM()], # Custom converters for specific rerun entity paths # Normally all these would be specified in their respectative modules # Until this is implemented we have central overrides here @@ -106,23 +103,14 @@ def _go2_rerun_blueprint() -> Any: }, } - -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge - - with_vis = autoconnect( - _transports_base, - FoxgloveBridge.blueprint(shm_channels=["/color_image#sensor_msgs.Image"]), - ) -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode - - with_vis = autoconnect( - _transports_base, - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config), - ) -else: - with_vis = _transports_base +with_vis = autoconnect( + _transports_base, + vis_module( + global_config.viewer, + rerun_config=rerun_config, + foxglove_config={"shm_channels": ["/color_image#sensor_msgs.Image"]}, + ), +) unitree_go2_basic = ( autoconnect( diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py new file mode 100644 index 0000000000..de786f67e8 --- /dev/null +++ b/dimos/visualization/vis_module.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# Copyright 2025-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. + +"""Shared visualization module factory for all robot blueprints.""" + +from typing import Any + +from dimos.core.blueprints import Blueprint, autoconnect +from dimos.core.global_config import ViewerBackend +from dimos.protocol.pubsub.impl.lcmpubsub import LCM + + +def vis_module( + viewer_backend: ViewerBackend, + rerun_config: dict[str, Any] | None = None, + foxglove_config: dict[str, Any] | None = None, +) -> Blueprint: + """Create a visualization blueprint based on the selected viewer backend. + + Bundles the appropriate viewer module (Rerun or Foxglove) together with + the ``RerunWebSocketServer`` so that remote viewer connections (click, + teleop) work out of the box when using a Rerun backend. + + Example usage:: + + from dimos.core.global_config import global_config + viz = vis_module( + global_config.viewer, + rerun_config={ + "visual_override": { + "world/camera_info": lambda ci: ci.to_rerun(...), + }, + "static": { + "world/tf/base_link": lambda rr: [rr.Boxes3D(...)], + }, + }, + ) + """ + if foxglove_config is None: + foxglove_config = {} + if rerun_config is None: + rerun_config = {} + rerun_config = {**rerun_config} + rerun_config.setdefault("pubsubs", [LCM()]) + + match viewer_backend: + case "foxglove": + from dimos.robot.foxglove_bridge import FoxgloveBridge + + return autoconnect(FoxgloveBridge.blueprint(**foxglove_config)) + case "rerun" | "rerun-web" | "rerun-connect": + from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + + viewer_mode = _BACKEND_TO_MODE.get(viewer_backend, "native") + return autoconnect( + RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), + RerunWebSocketServer.blueprint(), + ) + case _: + return autoconnect() From fa94c2eac320365fd57cb48826b226b7c5eb1348 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 16:49:38 -0700 Subject: [PATCH 06/53] cleanup --- dimos/robot/all_blueprints.py | 1 + .../go2/blueprints/basic/unitree_go2_basic.py | 4 ++-- .../visualization/rerun/test_viewer_ws_e2e.py | 19 ++------------- .../rerun/test_websocket_server.py | 23 ------------------- dimos/visualization/rerun/websocket_server.py | 21 ++++++++--------- 5 files changed, 14 insertions(+), 54 deletions(-) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 1fe034fd29..b18f934d1f 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -157,6 +157,7 @@ "reid-module": "dimos.perception.detection.reid.module", "replanning-a-star-planner": "dimos.navigation.replanning_a_star.module", "rerun-bridge-module": "dimos.visualization.rerun.bridge", + "rerun-web-socket-server": "dimos.visualization.rerun.websocket_server", "ros-nav": "dimos.navigation.rosnav", "simple-phone-teleop": "dimos.teleop.phone.phone_extensions", "simulation-module": "dimos.simulation.manipulators.sim_module", diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 406454ecc9..282f813571 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -103,7 +103,7 @@ def _go2_rerun_blueprint() -> Any: }, } -with_vis = autoconnect( +_with_vis = autoconnect( _transports_base, vis_module( global_config.viewer, @@ -114,7 +114,7 @@ def _go2_rerun_blueprint() -> Any: unitree_go2_basic = ( autoconnect( - with_vis, + _with_vis, GO2Connection.blueprint(), WebsocketVisModule.blueprint(), ) diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index d7bac7b6f4..5275adb660 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -29,6 +29,7 @@ import asyncio import json +import os import subprocess import threading import time @@ -39,11 +40,6 @@ _E2E_PORT = 13032 -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - def _make_server(port: int = _E2E_PORT) -> RerunWebSocketServer: return RerunWebSocketServer(port=port) @@ -77,11 +73,6 @@ async def _run() -> None: asyncio.run(_run()) -# --------------------------------------------------------------------------- -# Protocol-level E2E tests (no GUI required) -# --------------------------------------------------------------------------- - - class TestViewerProtocolE2E: """Verify the full Python-server side of the viewer ↔ DimOS protocol. @@ -264,11 +255,6 @@ def _on_pt(pt: Any) -> None: assert xs == [1.0, 2.0], f"Unexpected xs: {xs}" -# --------------------------------------------------------------------------- -# Binary smoke test -# --------------------------------------------------------------------------- - - class TestViewerBinaryConnectMode: """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket client attempts to connect to our Python server.""" @@ -297,9 +283,8 @@ def _on_pt(pt: Any) -> None: f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", ], env={ + **os.environ, "DISPLAY": "", - "HOME": "/home/dimos", - "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin", }, stdout=subprocess.PIPE, stderr=subprocess.PIPE, diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index c894774679..73c6759eec 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -29,11 +29,6 @@ _TEST_PORT = 13031 -# --------------------------------------------------------------------------- -# MockViewerPublisher -# --------------------------------------------------------------------------- - - class MockViewerPublisher: """Python mirror of the Rust WsPublisher in dimos-viewer. @@ -55,10 +50,6 @@ def __init__(self, url: str) -> None: self._ws: Any = None self._loop: asyncio.AbstractEventLoop | None = None - # ------------------------------------------------------------------ - # Context-manager interface - # ------------------------------------------------------------------ - def __enter__(self) -> "MockViewerPublisher": self._loop = asyncio.new_event_loop() self._ws = self._loop.run_until_complete(self._connect()) @@ -75,10 +66,6 @@ async def _connect(self) -> Any: return await ws_client.connect(self._url) - # ------------------------------------------------------------------ - # Send helpers (mirror of Rust WsPublisher methods) - # ------------------------------------------------------------------ - def send_click( self, x: float, @@ -138,11 +125,6 @@ def _send(self, msg: dict[str, Any]) -> None: self._loop.run_until_complete(self._ws.send(json.dumps(msg))) -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - def _collect(received: list[Any], done: threading.Event) -> Any: """Return a callback that appends to *received* and signals *done*.""" @@ -176,11 +158,6 @@ async def _probe() -> None: raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - class TestRerunWebSocketServerStartup: def test_server_binds_port(self) -> None: """After start(), the server must be reachable on the configured port.""" diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index ba0c953bd8..b374c739f0 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -77,10 +77,7 @@ def __init__(self, **kwargs: Any) -> None: self._ws_loop: asyncio.AbstractEventLoop | None = None self._server_thread: threading.Thread | None = None self._stop_event: asyncio.Event | None = None - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ + self._server_ready = threading.Event() @rpc def start(self) -> None: @@ -95,6 +92,9 @@ def start(self) -> None: @rpc def stop(self) -> None: + # Wait briefly for the server thread to initialise _stop_event so we + # don't silently skip the shutdown signal (race with _serve()). + self._server_ready.wait(timeout=5.0) if ( self._ws_loop is not None and not self._ws_loop.is_closed() @@ -103,10 +103,6 @@ def stop(self) -> None: self._ws_loop.call_soon_threadsafe(self._stop_event.set) super().stop() - # ------------------------------------------------------------------ - # Server - # ------------------------------------------------------------------ - def _run_server(self) -> None: """Entry point for the background server thread.""" self._ws_loop = asyncio.new_event_loop() @@ -120,6 +116,7 @@ async def _serve(self) -> None: import websockets.asyncio.server as ws_server self._stop_event = asyncio.Event() + self._server_ready.set() async with ws_server.serve( self._handle_client, @@ -140,10 +137,6 @@ async def _handle_client(self, websocket: Any) -> None: except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") - # ------------------------------------------------------------------ - # Message dispatch - # ------------------------------------------------------------------ - def _dispatch(self, raw: str | bytes) -> None: try: msg = json.loads(raw) @@ -151,6 +144,10 @@ def _dispatch(self, raw: str | bytes) -> None: logger.warning(f"RerunWebSocketServer: ignoring non-JSON message: {raw!r}") return + if not isinstance(msg, dict): + logger.warning(f"RerunWebSocketServer: expected JSON object, got {type(msg).__name__}") + return + msg_type = msg.get("type") if msg_type == "click": From 42f2f3860c2f2634c1bfd54a2fd1c4c0f8da1c2f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 17:13:58 -0700 Subject: [PATCH 07/53] consolidate viewer usage --- dimos/hardware/sensors/camera/module.py | 4 +- .../lidar/fastlio2/fastlio_blueprints.py | 39 ++++++++++++------- .../sensors/lidar/livox/livox_blueprints.py | 4 +- dimos/manipulation/blueprints.py | 4 +- dimos/manipulation/grasping/demo_grasping.py | 4 +- .../demo_object_scene_registration.py | 4 +- .../drone/blueprints/basic/drone_basic.py | 15 +------ .../blueprints/perceptive/unitree_g1_shm.py | 10 ++--- dimos/teleop/quest/blueprints.py | 4 +- 9 files changed, 43 insertions(+), 45 deletions(-) diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index b8165658d9..9c5623d141 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -32,7 +32,7 @@ from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier from dimos.spec import perception -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module def default_transform() -> Transform: @@ -120,5 +120,5 @@ def stop(self) -> None: demo_camera = autoconnect( CameraModule.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ) diff --git a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py index f3de842b46..b39dd7bcec 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py +++ b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py @@ -15,36 +15,45 @@ from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.lidar.fastlio2.module import FastLio2 from dimos.mapping.voxels import VoxelGridMapper -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module voxel_size = 0.05 mid360_fastlio = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=-1), - RerunBridgeModule.blueprint( - visual_override={ - "world/lidar": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), - } + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/lidar": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), + }, + }, ), ).global_config(n_workers=2, robot_model="mid360_fastlio2") mid360_fastlio_voxels = autoconnect( FastLio2.blueprint(), VoxelGridMapper.blueprint(publish_interval=1.0, voxel_size=voxel_size, carve_columns=False), - RerunBridgeModule.blueprint( - visual_override={ - "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), - "world/lidar": None, - } + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), + "world/lidar": None, + }, + }, ), ).global_config(n_workers=3, robot_model="mid360_fastlio2_voxels") mid360_fastlio_voxels_native = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=3.0), - RerunBridgeModule.blueprint( - visual_override={ - "world/lidar": None, - "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), - } + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/lidar": None, + "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), + }, + }, ), ).global_config(n_workers=2, robot_model="mid360_fastlio2") diff --git a/dimos/hardware/sensors/lidar/livox/livox_blueprints.py b/dimos/hardware/sensors/lidar/livox/livox_blueprints.py index c8835b3e89..958af084e2 100644 --- a/dimos/hardware/sensors/lidar/livox/livox_blueprints.py +++ b/dimos/hardware/sensors/lidar/livox/livox_blueprints.py @@ -14,9 +14,9 @@ from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.lidar.livox.module import Mid360 -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module mid360 = autoconnect( Mid360.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ).global_config(n_workers=2, robot_model="mid360") diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 8110166042..0a437bed1a 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -45,7 +45,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.robot.foxglove_bridge import FoxgloveBridge # TODO: migrate to rerun +from dimos.visualization.vis_module import vis_module from dimos.utils.data import get_data @@ -407,7 +407,7 @@ def _make_piper_config( base_transform=_XARM_PERCEPTION_CAMERA_TRANSFORM, ), ObjectSceneRegistrationModule.blueprint(target_frame="world"), - FoxgloveBridge.blueprint(), # TODO: migrate to rerun + vis_module("foxglove"), ) .transports( { diff --git a/dimos/manipulation/grasping/demo_grasping.py b/dimos/manipulation/grasping/demo_grasping.py index a4eea21787..f1ce67709e 100644 --- a/dimos/manipulation/grasping/demo_grasping.py +++ b/dimos/manipulation/grasping/demo_grasping.py @@ -21,7 +21,7 @@ from dimos.manipulation.grasping.grasping import GraspingModule from dimos.perception.detection.detectors.yoloe import YoloePromptMode from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.robot.foxglove_bridge import FoxgloveBridge +from dimos.visualization.vis_module import vis_module camera_module = RealSenseCamera.blueprint(enable_pointcloud=False) @@ -43,6 +43,6 @@ ("/tmp", "/tmp", "rw") ], # Grasp visualization debug standalone: python -m dimos.manipulation.grasping.visualize_grasps ), - FoxgloveBridge.blueprint(), + vis_module("foxglove"), Agent.blueprint(), ).global_config(viewer="foxglove") diff --git a/dimos/perception/demo_object_scene_registration.py b/dimos/perception/demo_object_scene_registration.py index 55b26f385a..13fb26cbb5 100644 --- a/dimos/perception/demo_object_scene_registration.py +++ b/dimos/perception/demo_object_scene_registration.py @@ -19,7 +19,7 @@ from dimos.hardware.sensors.camera.zed.compat import ZEDCamera from dimos.perception.detection.detectors.yoloe import YoloePromptMode from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.robot.foxglove_bridge import FoxgloveBridge +from dimos.visualization.vis_module import vis_module camera_choice = "zed" @@ -33,6 +33,6 @@ demo_object_scene_registration = autoconnect( camera_module, ObjectSceneRegistrationModule.blueprint(target_frame="world", prompt_mode=YoloePromptMode.LRPC), - FoxgloveBridge.blueprint(), + vis_module("foxglove"), Agent.blueprint(), ).global_config(viewer="foxglove") diff --git a/dimos/robot/drone/blueprints/basic/drone_basic.py b/dimos/robot/drone/blueprints/basic/drone_basic.py index fbe6621ae1..c99c273cc2 100644 --- a/dimos/robot/drone/blueprints/basic/drone_basic.py +++ b/dimos/robot/drone/blueprints/basic/drone_basic.py @@ -20,9 +20,9 @@ from dimos.core.blueprints import autoconnect from dimos.core.global_config import global_config -from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.robot.drone.camera_module import DroneCameraModule from dimos.robot.drone.connection_module import DroneConnectionModule +from dimos.visualization.vis_module import vis_module from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule @@ -60,23 +60,12 @@ def _drone_rerun_blueprint() -> Any: _rerun_config = { "blueprint": _drone_rerun_blueprint, - "pubsubs": [LCM()], "static": { "world/tf/base_link": _static_drone_body, }, } -# Conditional visualization -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge - - _vis = FoxgloveBridge.blueprint() -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode - - _vis = RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **_rerun_config) -else: - _vis = autoconnect() +_vis = vis_module(global_config.viewer, rerun_config=_rerun_config) # Determine connection string based on replay flag connection_string = "udp:0.0.0.0:14550" diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py index 5b127fb697..9efe400895 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py @@ -17,10 +17,11 @@ from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE from dimos.core.blueprints import autoconnect +from dimos.core.global_config import global_config from dimos.core.transport import pSHMTransport from dimos.msgs.sensor_msgs.Image import Image -from dimos.robot.foxglove_bridge import FoxgloveBridge from dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1 import unitree_g1 +from dimos.visualization.vis_module import vis_module unitree_g1_shm = autoconnect( unitree_g1.transports( @@ -30,10 +31,9 @@ ), } ), - FoxgloveBridge.blueprint( - shm_channels=[ - "/color_image#sensor_msgs.Image", - ] + vis_module( + global_config.viewer, + foxglove_config={"shm_channels": ["/color_image#sensor_msgs.Image"]}, ), ) diff --git a/dimos/teleop/quest/blueprints.py b/dimos/teleop/quest/blueprints.py index 6855ab62ca..9a044673b8 100644 --- a/dimos/teleop/quest/blueprints.py +++ b/dimos/teleop/quest/blueprints.py @@ -25,12 +25,12 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.teleop.quest.quest_extensions import ArmTeleopModule from dimos.teleop.quest.quest_types import Buttons -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module # Arm teleop with press-and-hold engage (has rerun viz) teleop_quest_rerun = autoconnect( ArmTeleopModule.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ).transports( { ("left_controller_output", PoseStamped): LCMTransport("/teleop/left_delta", PoseStamped), From 23d1d8887e1d7afd0bb7c8168efcf8a8cdb193e0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 17:20:16 -0700 Subject: [PATCH 08/53] consolidate WebsocketVisModule --- .../drone/blueprints/basic/drone_basic.py | 2 -- .../primitive/uintree_g1_primitive_no_nav.py | 3 --- .../go2/blueprints/basic/unitree_go2_basic.py | 2 -- .../go2/blueprints/basic/unitree_go2_fleet.py | 6 ++--- dimos/visualization/vis_module.py | 24 +++++++++++++++---- 5 files changed, 21 insertions(+), 16 deletions(-) diff --git a/dimos/robot/drone/blueprints/basic/drone_basic.py b/dimos/robot/drone/blueprints/basic/drone_basic.py index c99c273cc2..c60483cb0a 100644 --- a/dimos/robot/drone/blueprints/basic/drone_basic.py +++ b/dimos/robot/drone/blueprints/basic/drone_basic.py @@ -23,7 +23,6 @@ from dimos.robot.drone.camera_module import DroneCameraModule from dimos.robot.drone.connection_module import DroneConnectionModule from dimos.visualization.vis_module import vis_module -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule def _static_drone_body(rr: Any) -> list[Any]: @@ -81,7 +80,6 @@ def _drone_rerun_blueprint() -> Any: outdoor=False, ), DroneCameraModule.blueprint(camera_intrinsics=[1000.0, 1000.0, 960.0, 540.0]), - WebsocketVisModule.blueprint(), ) __all__ = [ diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index 2228dbfd66..220caff949 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -41,7 +41,6 @@ WavefrontFrontierExplorer, ) from dimos.visualization.vis_module import vis_module -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule def _convert_camera_info(camera_info: Any) -> Any: @@ -135,8 +134,6 @@ def _create_webcam() -> Webcam: VoxelGridMapper.blueprint(voxel_size=0.1), CostMapper.blueprint(), WavefrontFrontierExplorer.blueprint(), - # Visualization - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_g1") .transports( diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 282f813571..1e0f32d25c 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -25,7 +25,6 @@ from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.robot.unitree.go2.connection import GO2Connection from dimos.visualization.vis_module import vis_module -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule # Mac has some issue with high bandwidth UDP, so we use pSHMTransport for color_image # actually we can use pSHMTransport for all platforms, and for all streams @@ -116,7 +115,6 @@ def _go2_rerun_blueprint() -> Any: autoconnect( _with_vis, GO2Connection.blueprint(), - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py index 1c55f3e93c..0468cad40d 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py @@ -22,15 +22,13 @@ from dimos.core.blueprints import autoconnect from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator -from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import with_vis +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import _with_vis from dimos.robot.unitree.go2.fleet_connection import Go2FleetConnection -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule unitree_go2_fleet = ( autoconnect( - with_vis, + _with_vis, Go2FleetConnection.blueprint(), - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index de786f67e8..688a6efb5b 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -30,8 +30,8 @@ def vis_module( """Create a visualization blueprint based on the selected viewer backend. Bundles the appropriate viewer module (Rerun or Foxglove) together with - the ``RerunWebSocketServer`` so that remote viewer connections (click, - teleop) work out of the box when using a Rerun backend. + the ``WebsocketVisModule`` and ``RerunWebSocketServer`` so that the web + dashboard and remote viewer connections work out of the box. Example usage:: @@ -48,6 +48,8 @@ def vis_module( }, ) """ + from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule + if foxglove_config is None: foxglove_config = {} if rerun_config is None: @@ -59,8 +61,11 @@ def vis_module( case "foxglove": from dimos.robot.foxglove_bridge import FoxgloveBridge - return autoconnect(FoxgloveBridge.blueprint(**foxglove_config)) - case "rerun" | "rerun-web" | "rerun-connect": + return autoconnect( + FoxgloveBridge.blueprint(**foxglove_config), + WebsocketVisModule.blueprint(), + ) + case "rerun" | "rerun-web": from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule from dimos.visualization.rerun.websocket_server import RerunWebSocketServer @@ -68,6 +73,15 @@ def vis_module( return autoconnect( RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), RerunWebSocketServer.blueprint(), + WebsocketVisModule.blueprint(), + ) + case "rerun-connect": + from dimos.visualization.rerun.bridge import RerunBridgeModule + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + + return autoconnect( + RerunBridgeModule.blueprint(viewer_mode="connect", **rerun_config), + RerunWebSocketServer.blueprint(), ) case _: - return autoconnect() + return autoconnect(WebsocketVisModule.blueprint()) From 27e66be7a32301dd17d8edaaa4aa27d65574ed8b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 19:28:26 -0700 Subject: [PATCH 09/53] fix: address PR review - server ready race, path filter, skip guard - Move _server_ready.set() inside ws_server.serve() context so stop() waits for the port to actually bind before sending shutdown signal - Add /ws path filter to reject non-viewer WebSocket connections - Add pytest.mark.skipif for dimos-viewer binary test in CI - Fix import ordering in manipulation/blueprints.py --- dimos/manipulation/blueprints.py | 2 +- .../visualization/rerun/test_viewer_ws_e2e.py | 18 +++++++++++++++--- dimos/visualization/rerun/websocket_server.py | 5 ++++- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 0a437bed1a..aaad1c3525 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -45,8 +45,8 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.visualization.vis_module import vis_module from dimos.utils.data import get_data +from dimos.visualization.vis_module import vis_module def _make_base_pose( diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index 5275adb660..80c4743e61 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -30,11 +30,14 @@ import asyncio import json import os +import shutil import subprocess import threading import time from typing import Any +import pytest + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer _E2E_PORT = 13032 @@ -259,6 +262,13 @@ class TestViewerBinaryConnectMode: """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket client attempts to connect to our Python server.""" + @pytest.mark.skipif( + shutil.which("dimos-viewer") is None + or "--connect" not in subprocess.run( + ["dimos-viewer", "--help"], capture_output=True, text=True + ).stdout, + reason="dimos-viewer binary not installed or does not support --connect", + ) def test_viewer_ws_client_connects(self) -> None: """dimos-viewer --connect starts and its WS client connects to our server.""" server = _make_server() @@ -307,11 +317,13 @@ def _on_pt(pt: Any) -> None: except subprocess.TimeoutExpired: proc.kill() + stdout = proc.stdout.read().decode(errors="replace") if proc.stdout else "" stderr = proc.stderr.read().decode(errors="replace") if proc.stderr else "" server.stop() # The viewer should log that it is connecting to our WS URL. - # Even without a display, the log output appears before the GUI loop starts. - assert "ws://127.0.0.1" in stderr or proc.returncode is not None, ( - f"Viewer did not attempt WS connection. stderr:\n{stderr}" + # Check both stdout and stderr since log output destination varies. + combined = stdout + stderr + assert f"ws://127.0.0.1:{_E2E_PORT}" in combined, ( + f"Viewer did not attempt WS connection.\nstdout:\n{stdout}\nstderr:\n{stderr}" ) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index b374c739f0..16a292ca87 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -116,19 +116,22 @@ async def _serve(self) -> None: import websockets.asyncio.server as ws_server self._stop_event = asyncio.Event() - self._server_ready.set() async with ws_server.serve( self._handle_client, host=self.config.host, port=self.config.port, ): + self._server_ready.set() logger.info( f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" ) await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: + if hasattr(websocket, "request") and websocket.request.path != "/ws": + await websocket.close(1008, "Not Found") + return addr = websocket.remote_address logger.info(f"RerunWebSocketServer: viewer connected from {addr}") try: From fe84b4db10a8408173d000b0b64760ba5ba953a5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 19:31:14 -0700 Subject: [PATCH 10/53] fix: set explicit ping interval/timeout on WebSocket server The default websockets ping_interval=20s + ping_timeout=20s was too aggressive. Increase both to 30s to give the viewer more time to respond, especially during brief network hiccups. --- dimos/visualization/rerun/websocket_server.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 16a292ca87..b12307c11a 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -121,6 +121,10 @@ async def _serve(self) -> None: self._handle_client, host=self.config.host, port=self.config.port, + # Ping every 30 s, allow 30 s for pong — generous enough to + # survive brief network hiccups while still detecting dead clients. + ping_interval=30, + ping_timeout=30, ): self._server_ready.set() logger.info( From bbffeee7bee8899c1a1257ad1c56d9b878042bab Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:19:41 -0700 Subject: [PATCH 11/53] switch to websocket --- .../visualization/rerun/test_viewer_ws_e2e.py | 332 ++++++++++++++++ .../rerun/test_websocket_server.py | 374 ++++++++++++++++++ dimos/visualization/rerun/websocket_server.py | 173 ++++++++ 3 files changed, 879 insertions(+) create mode 100644 dimos/visualization/rerun/test_viewer_ws_e2e.py create mode 100644 dimos/visualization/rerun/test_websocket_server.py create mode 100644 dimos/visualization/rerun/websocket_server.py diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py new file mode 100644 index 0000000000..266e16cc68 --- /dev/null +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -0,0 +1,332 @@ +# Copyright 2025-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. + +"""End-to-end test: dimos-viewer (headless) → WebSocket → RerunWebSocketServer. + +dimos-viewer is started in ``--connect`` mode so it initialises its WebSocket +client. The viewer needs a gRPC proxy to connect to; we give it a non-existent +one so the viewer starts up anyway but produces no visualisation. The important +part is that the WebSocket client inside the viewer tries to connect to +``ws://127.0.0.1:/ws``. + +Because the viewer is a native GUI application it cannot run headlessly in CI +without a display. This test therefore verifies the connection at the protocol +level by using the ``RerunWebSocketServer`` module directly as the server and +injecting synthetic JSON messages that mimic what the viewer would send once a +user clicks in the 3D viewport. +""" + +import asyncio +import json +import subprocess +import threading +import time +from typing import Any + +import pytest + +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +_E2E_PORT = 13032 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_server(port: int = _E2E_PORT) -> RerunWebSocketServer: + return RerunWebSocketServer(port=port) + + +def _wait_for_server(port: int, timeout: float = 5.0) -> None: + import websockets.asyncio.client as ws_client + + async def _probe() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): + pass + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + asyncio.run(_probe()) + return + except Exception: + time.sleep(0.05) + raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") + + +def _send_messages(port: int, messages: list[dict[str, Any]], *, delay: float = 0.05) -> None: + import websockets.asyncio.client as ws_client + + async def _run() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws") as ws: + for msg in messages: + await ws.send(json.dumps(msg)) + await asyncio.sleep(delay) + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# Protocol-level E2E tests (no GUI required) +# --------------------------------------------------------------------------- + + +class TestViewerProtocolE2E: + """Verify the full Python-server side of the viewer ↔ DimOS protocol. + + These tests use the ``RerunWebSocketServer`` as the server and a dummy + WebSocket client (playing the role of dimos-viewer) to inject messages. + They confirm every message type is correctly routed and that only click + messages produce stream publishes. + """ + + def test_viewer_click_reaches_stream(self): + """A viewer click message received over WebSocket publishes PointStamped.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + done.set() + + server.clicked_point.subscribe(_on_pt) + + _send_messages( + _E2E_PORT, + [ + { + "type": "click", + "x": 10.0, + "y": 20.0, + "z": 0.5, + "entity_path": "/world/robot", + "timestamp_ms": 42000, + } + ], + ) + + done.wait(timeout=3.0) + server.stop() + + assert len(received) == 1 + pt = received[0] + assert abs(pt.x - 10.0) < 1e-9 + assert abs(pt.y - 20.0) < 1e-9 + assert abs(pt.z - 0.5) < 1e-9 + assert pt.frame_id == "/world/robot" + assert abs(pt.ts - 42.0) < 1e-6 + + def test_viewer_keyboard_twist_no_publish(self): + """Twist messages from keyboard control do not publish clicked_point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + server.clicked_point.subscribe(received.append) + + _send_messages( + _E2E_PORT, + [ + { + "type": "twist", + "linear_x": 0.5, + "linear_y": 0.0, + "linear_z": 0.0, + "angular_x": 0.0, + "angular_y": 0.0, + "angular_z": 0.8, + } + ], + ) + + server.stop() + assert received == [] + + def test_viewer_stop_no_publish(self): + """Stop messages do not publish clicked_point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + server.clicked_point.subscribe(received.append) + + _send_messages(_E2E_PORT, [{"type": "stop"}]) + + server.stop() + assert received == [] + + def test_full_viewer_session_sequence(self): + """Realistic session: connect, heartbeats, click, WASD, stop → one point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + done.set() + + server.clicked_point.subscribe(_on_pt) + + _send_messages( + _E2E_PORT, + [ + # Initial heartbeats (viewer connects and starts 1 Hz heartbeat) + {"type": "heartbeat", "timestamp_ms": 1000}, + {"type": "heartbeat", "timestamp_ms": 2000}, + # User clicks a point in the 3D viewport + { + "type": "click", + "x": 3.14, + "y": 2.71, + "z": 1.41, + "entity_path": "/world", + "timestamp_ms": 3000, + }, + # User presses W (forward) + { + "type": "twist", + "linear_x": 0.5, + "linear_y": 0.0, + "linear_z": 0.0, + "angular_x": 0.0, + "angular_y": 0.0, + "angular_z": 0.0, + }, + # User releases W + {"type": "stop"}, + # Another heartbeat + {"type": "heartbeat", "timestamp_ms": 4000}, + ], + delay=0.2, + ) + + done.wait(timeout=3.0) + server.stop() + + assert len(received) == 1, f"Expected exactly 1 click, got {len(received)}" + pt = received[0] + assert abs(pt.x - 3.14) < 1e-9 + assert abs(pt.y - 2.71) < 1e-9 + assert abs(pt.z - 1.41) < 1e-9 + + def test_reconnect_after_disconnect(self): + """Server keeps accepting new connections after a client disconnects.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + all_done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + if len(received) >= 2: + all_done.set() + + server.clicked_point.subscribe(_on_pt) + + # First connection — send one click and disconnect + _send_messages( + _E2E_PORT, + [{"type": "click", "x": 1.0, "y": 0.0, "z": 0.0, "entity_path": "", "timestamp_ms": 0}], + ) + + # Second connection (simulating viewer reconnect) — send another click + _send_messages( + _E2E_PORT, + [{"type": "click", "x": 2.0, "y": 0.0, "z": 0.0, "entity_path": "", "timestamp_ms": 0}], + ) + + all_done.wait(timeout=5.0) + server.stop() + + xs = sorted(pt.x for pt in received) + assert xs == [1.0, 2.0], f"Unexpected xs: {xs}" + + +# --------------------------------------------------------------------------- +# Binary smoke test +# --------------------------------------------------------------------------- + + +class TestViewerBinaryConnectMode: + """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket + client attempts to connect to our Python server.""" + + def test_viewer_ws_client_connects(self): + """dimos-viewer --connect starts and its WS client connects to our server.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + connected = threading.Event() + received: list[Any] = [] + + def _on_pt(pt: Any) -> None: + received.append(pt) + + server.clicked_point.subscribe(_on_pt) + + # Start dimos-viewer in --connect mode, pointing it at a non-existent gRPC + # proxy (it will fail to stream data, but that's fine) and at our WS server. + # Use DISPLAY="" to prevent it from opening a window (it will exit quickly + # without a display, but the WebSocket connection happens before the GUI loop). + proc = subprocess.Popen( + [ + "dimos-viewer", + "--connect", + f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", + ], + env={"DISPLAY": "", "HOME": "/home/dimos", "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin"}, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + # Give the viewer up to 5 s to connect its WebSocket client to our server. + # We detect the connection by waiting for the server to accept a client. + deadline = time.monotonic() + 5.0 + viewer_connected = False + while time.monotonic() < deadline: + # Check if any connection was established by sending a message and + # verifying the viewer is still running. + if proc.poll() is not None: + # Viewer exited (expected without a display) — check if it connected first. + break + time.sleep(0.1) + + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + + stderr = proc.stderr.read().decode(errors="replace") if proc.stderr else "" + server.stop() + + # The viewer should log that it is connecting to our WS URL. + # Even without a display, the log output appears before the GUI loop starts. + assert "ws://127.0.0.1" in stderr or proc.returncode is not None, ( + f"Viewer did not attempt WS connection. stderr:\n{stderr}" + ) diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py new file mode 100644 index 0000000000..e1dc08ee23 --- /dev/null +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -0,0 +1,374 @@ +# Copyright 2025-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. + +"""Tests for RerunWebSocketServer. + +Uses ``MockViewerPublisher`` to simulate dimos-viewer sending events, matching +the exact JSON protocol used by the Rust ``WsPublisher`` in the viewer. +""" + +import asyncio +import json +import threading +import time +from typing import Any + +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +_TEST_PORT = 13031 + + +# --------------------------------------------------------------------------- +# MockViewerPublisher +# --------------------------------------------------------------------------- + + +class MockViewerPublisher: + """Python mirror of the Rust WsPublisher in dimos-viewer. + + Connects to a running ``RerunWebSocketServer`` and exposes the same + ``send_click`` / ``send_twist`` / ``send_stop`` / ``send_heartbeat`` + API that the real viewer uses. Useful for unit tests that need to + exercise the server without a real viewer binary. + + Usage:: + + with MockViewerPublisher("ws://127.0.0.1:13031/ws") as pub: + pub.send_click(1.0, 2.0, 0.0, "/world", timestamp_ms=1000) + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.send_stop() + """ + + def __init__(self, url: str) -> None: + self._url = url + self._ws: Any = None + self._loop: asyncio.AbstractEventLoop | None = None + + # ------------------------------------------------------------------ + # Context-manager interface + # ------------------------------------------------------------------ + + def __enter__(self) -> "MockViewerPublisher": + self._loop = asyncio.new_event_loop() + self._ws = self._loop.run_until_complete(self._connect()) + return self + + def __exit__(self, *_: Any) -> None: + if self._ws is not None and self._loop is not None: + self._loop.run_until_complete(self._ws.close()) + if self._loop is not None: + self._loop.close() + + async def _connect(self) -> Any: + import websockets.asyncio.client as ws_client + return await ws_client.connect(self._url) + + # ------------------------------------------------------------------ + # Send helpers (mirror of Rust WsPublisher methods) + # ------------------------------------------------------------------ + + def send_click( + self, + x: float, + y: float, + z: float, + entity_path: str = "", + timestamp_ms: int = 0, + ) -> None: + """Send a click event — matches viewer SelectionChange handler output.""" + self._send( + { + "type": "click", + "x": x, + "y": y, + "z": z, + "entity_path": entity_path, + "timestamp_ms": timestamp_ms, + } + ) + + def send_twist( + self, + linear_x: float, + linear_y: float, + linear_z: float, + angular_x: float, + angular_y: float, + angular_z: float, + ) -> None: + """Send a twist (WASD keyboard) event.""" + self._send( + { + "type": "twist", + "linear_x": linear_x, + "linear_y": linear_y, + "linear_z": linear_z, + "angular_x": angular_x, + "angular_y": angular_y, + "angular_z": angular_z, + } + ) + + def send_stop(self) -> None: + """Send a stop event (Space bar or key release).""" + self._send({"type": "stop"}) + + def send_heartbeat(self, timestamp_ms: int = 0) -> None: + """Send a heartbeat (1 Hz keepalive from viewer).""" + self._send({"type": "heartbeat", "timestamp_ms": timestamp_ms}) + + def flush(self, delay: float = 0.1) -> None: + """Wait briefly so the server processes queued messages.""" + time.sleep(delay) + + def _send(self, msg: dict[str, Any]) -> None: + assert self._loop is not None and self._ws is not None, "Not connected" + self._loop.run_until_complete(self._ws.send(json.dumps(msg))) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: + return RerunWebSocketServer(port=port) + + +def _wait_for_server(port: int, timeout: float = 3.0) -> None: + """Block until the WebSocket server accepts an upgrade handshake.""" + async def _probe() -> None: + import websockets.asyncio.client as ws_client + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): + pass + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + asyncio.run(_probe()) + return + except Exception: + time.sleep(0.05) + raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestRerunWebSocketServerStartup: + def test_server_binds_port(self) -> None: + """After start(), the server must be reachable on the configured port.""" + mod = _make_module() + mod.start() + try: + _wait_for_server(_TEST_PORT) + finally: + mod.stop() + + def test_stop_is_idempotent(self) -> None: + """Calling stop() twice must not raise.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + mod.stop() + mod.stop() + + +class TestClickMessages: + def test_click_publishes_point_stamped(self) -> None: + """A single click publishes one PointStamped with correct coords.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(1.5, 2.5, 0.0, "/world", timestamp_ms=1000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + pt = received[0] + assert abs(pt.x - 1.5) < 1e-9 + assert abs(pt.y - 2.5) < 1e-9 + assert abs(pt.z - 0.0) < 1e-9 + + def test_click_sets_frame_id_from_entity_path(self) -> None: + """entity_path is stored as frame_id on the published PointStamped.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(0.0, 0.0, 0.0, "/robot/base", timestamp_ms=2000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + assert received and received[0].frame_id == "/robot/base" + + def test_click_timestamp_converted_from_ms(self) -> None: + """timestamp_ms is converted to seconds on PointStamped.ts.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(0.0, 0.0, 0.0, "", timestamp_ms=5000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + assert received and abs(received[0].ts - 5.0) < 1e-6 + + def test_multiple_clicks_all_published(self) -> None: + """A burst of clicks all arrive on the stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + all_arrived = threading.Event() + + def _cb(pt: Any) -> None: + received.append(pt) + if len(received) >= 3: + all_arrived.set() + + mod.clicked_point.subscribe(_cb) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(1.0, 0.0, 0.0) + pub.send_click(2.0, 0.0, 0.0) + pub.send_click(3.0, 0.0, 0.0) + pub.flush() + + all_arrived.wait(timeout=3.0) + mod.stop() + + assert sorted(pt.x for pt in received) == [1.0, 2.0, 3.0] + + +class TestNonClickMessages: + def test_heartbeat_does_not_publish(self) -> None: + """Heartbeat messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_heartbeat(9999) + pub.flush() + + mod.stop() + assert received == [] + + def test_twist_does_not_publish_clicked_point(self) -> None: + """Twist messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.flush() + + mod.stop() + assert received == [] + + def test_stop_does_not_publish_clicked_point(self) -> None: + """Stop messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_stop() + pub.flush() + + mod.stop() + assert received == [] + + def test_invalid_json_does_not_crash(self) -> None: + """Malformed JSON is silently dropped; server stays alive.""" + import websockets.asyncio.client as ws_client + + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + async def _send_bad() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{_TEST_PORT}/ws") as ws: + await ws.send("this is not json {{") + await asyncio.sleep(0.1) + await ws.send(json.dumps({"type": "heartbeat", "timestamp_ms": 0})) + await asyncio.sleep(0.1) + + asyncio.run(_send_bad()) + mod.stop() + + def test_mixed_message_sequence(self) -> None: + """Realistic sequence: heartbeat → click → twist → stop publishes one point.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + # Subscribe before sending so we don't race against the click dispatch. + received: list[Any] = [] + done = threading.Event() + + def _cb(pt: Any) -> None: + received.append(pt) + done.set() + + mod.clicked_point.subscribe(_cb) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_heartbeat(1000) + pub.send_click(7.0, 8.0, 9.0, "/map", timestamp_ms=1100) + pub.send_twist(0.3, 0.0, 0.0, 0.0, 0.0, 0.2) + pub.send_stop() + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + assert abs(received[0].x - 7.0) < 1e-9 + assert abs(received[0].y - 8.0) < 1e-9 + assert abs(received[0].z - 9.0) < 1e-9 diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py new file mode 100644 index 0000000000..a04e2c4999 --- /dev/null +++ b/dimos/visualization/rerun/websocket_server.py @@ -0,0 +1,173 @@ +# Copyright 2025-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. + +"""WebSocket server module that receives events from dimos-viewer. + +When dimos-viewer is started with ``--connect``, LCM multicast is unavailable +across machines. The viewer falls back to sending click, twist, and stop events +as JSON over a WebSocket connection. This module acts as the server-side +counterpart: it listens for those connections and translates incoming messages +into DimOS stream publishes. + +Message format (newline-delimited JSON, ``"type"`` discriminant): + + {"type":"heartbeat","timestamp_ms":1234567890} + {"type":"click","x":1.0,"y":2.0,"z":3.0,"entity_path":"/world","timestamp_ms":1234567890} + {"type":"twist","linear_x":0.5,"linear_y":0.0,"linear_z":0.0, + "angular_x":0.0,"angular_y":0.0,"angular_z":0.8} + {"type":"stop"} +""" + +import asyncio +import json +import threading +from typing import Any + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out +from dimos.msgs.geometry_msgs.PointStamped import PointStamped +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class Config(ModuleConfig): + port: int = 3030 + + +class RerunWebSocketServer(Module[Config]): + """Receives dimos-viewer WebSocket events and publishes them as DimOS streams. + + The viewer connects to this module (not the other way around) when running + in ``--connect`` mode. Each click event is converted to a ``PointStamped`` + and published on the ``clicked_point`` stream so downstream modules (e.g. + ``ReplanningAStarPlanner``) can consume it without modification. + + Outputs: + clicked_point: 3-D world-space point from the most recent viewer click. + """ + + default_config = Config + + clicked_point: Out[PointStamped] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._ws_loop: asyncio.AbstractEventLoop | None = None + self._server_thread: threading.Thread | None = None + self._stop_event: asyncio.Event | None = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + @rpc + def start(self) -> None: + super().start() + self._server_thread = threading.Thread( + target=self._run_server, daemon=True, name="rerun-ws-server" + ) + self._server_thread.start() + logger.info(f"RerunWebSocketServer starting on ws://0.0.0.0:{self.config.port}/ws") + + @rpc + def stop(self) -> None: + if ( + self._ws_loop is not None + and not self._ws_loop.is_closed() + and self._stop_event is not None + ): + self._ws_loop.call_soon_threadsafe(self._stop_event.set) + super().stop() + + # ------------------------------------------------------------------ + # Server + # ------------------------------------------------------------------ + + def _run_server(self) -> None: + """Entry point for the background server thread.""" + self._ws_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._ws_loop) + try: + self._ws_loop.run_until_complete(self._serve()) + finally: + self._ws_loop.close() + + async def _serve(self) -> None: + import websockets.asyncio.server as ws_server + + self._stop_event = asyncio.Event() + + async with ws_server.serve( + self._handle_client, + host="0.0.0.0", + port=self.config.port, + ): + logger.info( + f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws" + ) + await self._stop_event.wait() + + async def _handle_client(self, websocket: Any) -> None: + addr = websocket.remote_address + logger.info(f"RerunWebSocketServer: viewer connected from {addr}") + try: + async for raw in websocket: + self._dispatch(raw) + except Exception as exc: + logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") + + # ------------------------------------------------------------------ + # Message dispatch + # ------------------------------------------------------------------ + + def _dispatch(self, raw: str | bytes) -> None: + try: + msg = json.loads(raw) + except json.JSONDecodeError: + logger.warning(f"RerunWebSocketServer: ignoring non-JSON message: {raw!r}") + return + + msg_type = msg.get("type") + + if msg_type == "click": + pt = PointStamped( + x=float(msg["x"]), + y=float(msg["y"]), + z=float(msg["z"]), + ts=float(msg.get("timestamp_ms", 0)) / 1000.0, + frame_id=str(msg.get("entity_path", "")), + ) + logger.debug(f"RerunWebSocketServer: click → {pt}") + self.clicked_point.publish(pt) + + elif msg_type == "twist": + # Twist messages are not yet wired to a stream; log for observability. + logger.debug( + "RerunWebSocketServer: twist lin=({linear_x},{linear_y},{linear_z}) " + "ang=({angular_x},{angular_y},{angular_z})".format(**msg) + ) + + elif msg_type == "stop": + logger.debug("RerunWebSocketServer: stop") + + elif msg_type == "heartbeat": + logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") + + else: + logger.warning(f"RerunWebSocketServer: unknown message type {msg_type!r}") + + +rerun_ws_server = RerunWebSocketServer.blueprint From f99b2d4ed4ae11e69c0d7fa48dd15090c1d571e5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:27:40 -0700 Subject: [PATCH 12/53] cleanup --- .../visualization/rerun/test_viewer_ws_e2e.py | 11 ++++---- .../rerun/test_websocket_server.py | 3 +++ dimos/visualization/rerun/websocket_server.py | 26 ++++++++++++++----- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index 266e16cc68..4026d1d346 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -34,8 +34,6 @@ import time from typing import Any -import pytest - from dimos.visualization.rerun.websocket_server import RerunWebSocketServer _E2E_PORT = 13032 @@ -281,7 +279,7 @@ def test_viewer_ws_client_connects(self): server.start() _wait_for_server(_E2E_PORT) - connected = threading.Event() + threading.Event() received: list[Any] = [] def _on_pt(pt: Any) -> None: @@ -299,7 +297,11 @@ def _on_pt(pt: Any) -> None: "--connect", f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", ], - env={"DISPLAY": "", "HOME": "/home/dimos", "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin"}, + env={ + "DISPLAY": "", + "HOME": "/home/dimos", + "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin", + }, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) @@ -307,7 +309,6 @@ def _on_pt(pt: Any) -> None: # Give the viewer up to 5 s to connect its WebSocket client to our server. # We detect the connection by waiting for the server to accept a client. deadline = time.monotonic() + 5.0 - viewer_connected = False while time.monotonic() < deadline: # Check if any connection was established by sending a message and # verifying the viewer is still running. diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index e1dc08ee23..d0bd986d91 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -72,6 +72,7 @@ def __exit__(self, *_: Any) -> None: async def _connect(self) -> Any: import websockets.asyncio.client as ws_client + return await ws_client.connect(self._url) # ------------------------------------------------------------------ @@ -148,8 +149,10 @@ def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: def _wait_for_server(port: int, timeout: float = 3.0) -> None: """Block until the WebSocket server accepts an upgrade handshake.""" + async def _probe() -> None: import websockets.asyncio.client as ws_client + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): pass diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index a04e2c4999..70b6468408 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -38,6 +38,8 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.msgs.geometry_msgs.PointStamped import PointStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -57,11 +59,13 @@ class RerunWebSocketServer(Module[Config]): Outputs: clicked_point: 3-D world-space point from the most recent viewer click. + tele_cmd_vel: Twist velocity commands from keyboard teleop, including stop events. """ default_config = Config clicked_point: Out[PointStamped] + tele_cmd_vel: Out[Twist] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -115,9 +119,7 @@ async def _serve(self) -> None: host="0.0.0.0", port=self.config.port, ): - logger.info( - f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws" - ) + logger.info(f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws") await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -154,14 +156,24 @@ def _dispatch(self, raw: str | bytes) -> None: self.clicked_point.publish(pt) elif msg_type == "twist": - # Twist messages are not yet wired to a stream; log for observability. - logger.debug( - "RerunWebSocketServer: twist lin=({linear_x},{linear_y},{linear_z}) " - "ang=({angular_x},{angular_y},{angular_z})".format(**msg) + twist = Twist( + linear=Vector3( + float(msg.get("linear_x", 0)), + float(msg.get("linear_y", 0)), + float(msg.get("linear_z", 0)), + ), + angular=Vector3( + float(msg.get("angular_x", 0)), + float(msg.get("angular_y", 0)), + float(msg.get("angular_z", 0)), + ), ) + logger.debug(f"RerunWebSocketServer: twist → {twist}") + self.tele_cmd_vel.publish(twist) elif msg_type == "stop": logger.debug("RerunWebSocketServer: stop") + self.tele_cmd_vel.publish(Twist.zero()) elif msg_type == "heartbeat": logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") From 26273d737f59bf7b6bcf195183ad6504738efc4f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:33:47 -0700 Subject: [PATCH 13/53] improvements --- .../visualization/rerun/test_viewer_ws_e2e.py | 13 ++-- .../rerun/test_websocket_server.py | 59 ++++++++++++++++++- dimos/visualization/rerun/websocket_server.py | 13 ++-- 3 files changed, 71 insertions(+), 14 deletions(-) diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index 4026d1d346..d7bac7b6f4 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -91,7 +91,7 @@ class TestViewerProtocolE2E: messages produce stream publishes. """ - def test_viewer_click_reaches_stream(self): + def test_viewer_click_reaches_stream(self) -> None: """A viewer click message received over WebSocket publishes PointStamped.""" server = _make_server() server.start() @@ -131,7 +131,7 @@ def _on_pt(pt: Any) -> None: assert pt.frame_id == "/world/robot" assert abs(pt.ts - 42.0) < 1e-6 - def test_viewer_keyboard_twist_no_publish(self): + def test_viewer_keyboard_twist_no_publish(self) -> None: """Twist messages from keyboard control do not publish clicked_point.""" server = _make_server() server.start() @@ -158,7 +158,7 @@ def test_viewer_keyboard_twist_no_publish(self): server.stop() assert received == [] - def test_viewer_stop_no_publish(self): + def test_viewer_stop_no_publish(self) -> None: """Stop messages do not publish clicked_point.""" server = _make_server() server.start() @@ -172,7 +172,7 @@ def test_viewer_stop_no_publish(self): server.stop() assert received == [] - def test_full_viewer_session_sequence(self): + def test_full_viewer_session_sequence(self) -> None: """Realistic session: connect, heartbeats, click, WASD, stop → one point.""" server = _make_server() server.start() @@ -229,7 +229,7 @@ def _on_pt(pt: Any) -> None: assert abs(pt.y - 2.71) < 1e-9 assert abs(pt.z - 1.41) < 1e-9 - def test_reconnect_after_disconnect(self): + def test_reconnect_after_disconnect(self) -> None: """Server keeps accepting new connections after a client disconnects.""" server = _make_server() server.start() @@ -273,13 +273,12 @@ class TestViewerBinaryConnectMode: """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket client attempts to connect to our Python server.""" - def test_viewer_ws_client_connects(self): + def test_viewer_ws_client_connects(self) -> None: """dimos-viewer --connect starts and its WS client connects to our server.""" server = _make_server() server.start() _wait_for_server(_E2E_PORT) - threading.Event() received: list[Any] = [] def _on_pt(pt: Any) -> None: diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index d0bd986d91..c894774679 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -143,6 +143,16 @@ def _send(self, msg: dict[str, Any]) -> None: # --------------------------------------------------------------------------- +def _collect(received: list[Any], done: threading.Event) -> Any: + """Return a callback that appends to *received* and signals *done*.""" + + def _cb(msg: Any) -> None: + received.append(msg) + done.set() + + return _cb + + def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: return RerunWebSocketServer(port=port) @@ -199,7 +209,7 @@ def test_click_publishes_point_stamped(self) -> None: received: list[Any] = [] done = threading.Event() - mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + mod.clicked_point.subscribe(_collect(received, done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_click(1.5, 2.5, 0.0, "/world", timestamp_ms=1000) @@ -222,7 +232,7 @@ def test_click_sets_frame_id_from_entity_path(self) -> None: received: list[Any] = [] done = threading.Event() - mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + mod.clicked_point.subscribe(_collect(received, done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_click(0.0, 0.0, 0.0, "/robot/base", timestamp_ms=2000) @@ -240,7 +250,7 @@ def test_click_timestamp_converted_from_ms(self) -> None: received: list[Any] = [] done = threading.Event() - mod.clicked_point.subscribe(lambda pt: (received.append(pt), done.set())) + mod.clicked_point.subscribe(_collect(received, done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_click(0.0, 0.0, 0.0, "", timestamp_ms=5000) @@ -327,6 +337,49 @@ def test_stop_does_not_publish_clicked_point(self) -> None: mod.stop() assert received == [] + def test_twist_publishes_on_tele_cmd_vel(self) -> None: + """Twist messages publish a Twist on the tele_cmd_vel stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + tw = received[0] + assert abs(tw.linear.x - 0.5) < 1e-9 + assert abs(tw.angular.z - 0.8) < 1e-9 + + def test_stop_publishes_zero_twist_on_tele_cmd_vel(self) -> None: + """Stop messages publish a zero Twist on the tele_cmd_vel stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_stop() + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + tw = received[0] + assert tw.is_zero() + def test_invalid_json_does_not_crash(self) -> None: """Malformed JSON is silently dropped; server stays alive.""" import websockets.asyncio.client as ws_client diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 70b6468408..163bfcbf62 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -34,6 +34,8 @@ import threading from typing import Any +import websockets + from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out @@ -46,6 +48,9 @@ class Config(ModuleConfig): + # Intentionally binds 0.0.0.0 by default so the viewer can connect from + # any machine on the network (the typical robot deployment scenario). + host: str = "0.0.0.0" port: int = 3030 @@ -84,7 +89,7 @@ def start(self) -> None: target=self._run_server, daemon=True, name="rerun-ws-server" ) self._server_thread.start() - logger.info(f"RerunWebSocketServer starting on ws://0.0.0.0:{self.config.port}/ws") + logger.info(f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws") @rpc def stop(self) -> None: @@ -116,10 +121,10 @@ async def _serve(self) -> None: async with ws_server.serve( self._handle_client, - host="0.0.0.0", + host=self.config.host, port=self.config.port, ): - logger.info(f"RerunWebSocketServer listening on ws://0.0.0.0:{self.config.port}/ws") + logger.info(f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws") await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -128,7 +133,7 @@ async def _handle_client(self, websocket: Any) -> None: try: async for raw in websocket: self._dispatch(raw) - except Exception as exc: + except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") # ------------------------------------------------------------------ From 62fb365a704dd8b10e87ab356ebdbf424075bd05 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:46:25 -0700 Subject: [PATCH 14/53] fix: ruff formatting + consistent error handling in websocket_server --- dimos/visualization/rerun/websocket_server.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 163bfcbf62..ba0c953bd8 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -89,7 +89,9 @@ def start(self) -> None: target=self._run_server, daemon=True, name="rerun-ws-server" ) self._server_thread.start() - logger.info(f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws") + logger.info( + f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws" + ) @rpc def stop(self) -> None: @@ -124,7 +126,9 @@ async def _serve(self) -> None: host=self.config.host, port=self.config.port, ): - logger.info(f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws") + logger.info( + f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" + ) await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -151,9 +155,9 @@ def _dispatch(self, raw: str | bytes) -> None: if msg_type == "click": pt = PointStamped( - x=float(msg["x"]), - y=float(msg["y"]), - z=float(msg["z"]), + x=float(msg.get("x", 0)), + y=float(msg.get("y", 0)), + z=float(msg.get("z", 0)), ts=float(msg.get("timestamp_ms", 0)) / 1000.0, frame_id=str(msg.get("entity_path", "")), ) From 238b3de582af62ea03f4e362a903ce769658b619 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 16:43:17 -0700 Subject: [PATCH 15/53] make it easy to use --- .../primitive/uintree_g1_primitive_no_nav.py | 16 +--- .../go2/blueprints/basic/unitree_go2_basic.py | 30 +++----- dimos/visualization/vis_module.py | 73 +++++++++++++++++++ 3 files changed, 84 insertions(+), 35 deletions(-) create mode 100644 dimos/visualization/vis_module.py diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index c3da9521c5..2228dbfd66 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -40,7 +40,7 @@ from dimos.navigation.frontier_exploration.wavefront_frontier_goal_selector import ( WavefrontFrontierExplorer, ) -from dimos.protocol.pubsub.impl.lcmpubsub import LCM +from dimos.visualization.vis_module import vis_module from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule @@ -90,7 +90,6 @@ def _g1_rerun_blueprint() -> Any: rerun_config = { "blueprint": _g1_rerun_blueprint, - "pubsubs": [LCM()], "visual_override": { "world/camera_info": _convert_camera_info, "world/global_map": _convert_global_map, @@ -101,18 +100,7 @@ def _g1_rerun_blueprint() -> Any: }, } -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge - - _with_vis = autoconnect(FoxgloveBridge.blueprint()) -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode - - _with_vis = autoconnect( - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config) - ) -else: - _with_vis = autoconnect() +_with_vis = vis_module(global_config.viewer, rerun_config=rerun_config) def _create_webcam() -> Webcam: diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index a0d1e6a7ae..406454ecc9 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -22,9 +22,9 @@ from dimos.core.global_config import global_config from dimos.core.transport import pSHMTransport from dimos.msgs.sensor_msgs.Image import Image -from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.visualization.vis_module import vis_module from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule # Mac has some issue with high bandwidth UDP, so we use pSHMTransport for color_image @@ -87,9 +87,6 @@ def _go2_rerun_blueprint() -> Any: rerun_config = { "blueprint": _go2_rerun_blueprint, - # any pubsub that supports subscribe_all and topic that supports str(topic) - # is acceptable here - "pubsubs": [LCM()], # Custom converters for specific rerun entity paths # Normally all these would be specified in their respectative modules # Until this is implemented we have central overrides here @@ -106,23 +103,14 @@ def _go2_rerun_blueprint() -> Any: }, } - -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge - - with_vis = autoconnect( - _transports_base, - FoxgloveBridge.blueprint(shm_channels=["/color_image#sensor_msgs.Image"]), - ) -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode - - with_vis = autoconnect( - _transports_base, - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config), - ) -else: - with_vis = _transports_base +with_vis = autoconnect( + _transports_base, + vis_module( + global_config.viewer, + rerun_config=rerun_config, + foxglove_config={"shm_channels": ["/color_image#sensor_msgs.Image"]}, + ), +) unitree_go2_basic = ( autoconnect( diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py new file mode 100644 index 0000000000..de786f67e8 --- /dev/null +++ b/dimos/visualization/vis_module.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# Copyright 2025-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. + +"""Shared visualization module factory for all robot blueprints.""" + +from typing import Any + +from dimos.core.blueprints import Blueprint, autoconnect +from dimos.core.global_config import ViewerBackend +from dimos.protocol.pubsub.impl.lcmpubsub import LCM + + +def vis_module( + viewer_backend: ViewerBackend, + rerun_config: dict[str, Any] | None = None, + foxglove_config: dict[str, Any] | None = None, +) -> Blueprint: + """Create a visualization blueprint based on the selected viewer backend. + + Bundles the appropriate viewer module (Rerun or Foxglove) together with + the ``RerunWebSocketServer`` so that remote viewer connections (click, + teleop) work out of the box when using a Rerun backend. + + Example usage:: + + from dimos.core.global_config import global_config + viz = vis_module( + global_config.viewer, + rerun_config={ + "visual_override": { + "world/camera_info": lambda ci: ci.to_rerun(...), + }, + "static": { + "world/tf/base_link": lambda rr: [rr.Boxes3D(...)], + }, + }, + ) + """ + if foxglove_config is None: + foxglove_config = {} + if rerun_config is None: + rerun_config = {} + rerun_config = {**rerun_config} + rerun_config.setdefault("pubsubs", [LCM()]) + + match viewer_backend: + case "foxglove": + from dimos.robot.foxglove_bridge import FoxgloveBridge + + return autoconnect(FoxgloveBridge.blueprint(**foxglove_config)) + case "rerun" | "rerun-web" | "rerun-connect": + from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + + viewer_mode = _BACKEND_TO_MODE.get(viewer_backend, "native") + return autoconnect( + RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), + RerunWebSocketServer.blueprint(), + ) + case _: + return autoconnect() From 5c14fe2495896ca2bda44e886d7279408c055ddc Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 16:49:38 -0700 Subject: [PATCH 16/53] cleanup --- dimos/robot/all_blueprints.py | 1 + .../go2/blueprints/basic/unitree_go2_basic.py | 4 ++-- .../visualization/rerun/test_viewer_ws_e2e.py | 19 ++------------- .../rerun/test_websocket_server.py | 23 ------------------- dimos/visualization/rerun/websocket_server.py | 21 ++++++++--------- 5 files changed, 14 insertions(+), 54 deletions(-) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 5910093d61..44bfa8e280 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -157,6 +157,7 @@ "reid-module": "dimos.perception.detection.reid.module", "replanning-a-star-planner": "dimos.navigation.replanning_a_star.module", "rerun-bridge-module": "dimos.visualization.rerun.bridge", + "rerun-web-socket-server": "dimos.visualization.rerun.websocket_server", "ros-nav": "dimos.navigation.rosnav", "simple-phone-teleop": "dimos.teleop.phone.phone_extensions", "spatial-memory": "dimos.perception.spatial_perception", diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 406454ecc9..282f813571 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -103,7 +103,7 @@ def _go2_rerun_blueprint() -> Any: }, } -with_vis = autoconnect( +_with_vis = autoconnect( _transports_base, vis_module( global_config.viewer, @@ -114,7 +114,7 @@ def _go2_rerun_blueprint() -> Any: unitree_go2_basic = ( autoconnect( - with_vis, + _with_vis, GO2Connection.blueprint(), WebsocketVisModule.blueprint(), ) diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index d7bac7b6f4..5275adb660 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -29,6 +29,7 @@ import asyncio import json +import os import subprocess import threading import time @@ -39,11 +40,6 @@ _E2E_PORT = 13032 -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - def _make_server(port: int = _E2E_PORT) -> RerunWebSocketServer: return RerunWebSocketServer(port=port) @@ -77,11 +73,6 @@ async def _run() -> None: asyncio.run(_run()) -# --------------------------------------------------------------------------- -# Protocol-level E2E tests (no GUI required) -# --------------------------------------------------------------------------- - - class TestViewerProtocolE2E: """Verify the full Python-server side of the viewer ↔ DimOS protocol. @@ -264,11 +255,6 @@ def _on_pt(pt: Any) -> None: assert xs == [1.0, 2.0], f"Unexpected xs: {xs}" -# --------------------------------------------------------------------------- -# Binary smoke test -# --------------------------------------------------------------------------- - - class TestViewerBinaryConnectMode: """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket client attempts to connect to our Python server.""" @@ -297,9 +283,8 @@ def _on_pt(pt: Any) -> None: f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", ], env={ + **os.environ, "DISPLAY": "", - "HOME": "/home/dimos", - "PATH": "/home/dimos/.cargo/bin:/usr/bin:/bin", }, stdout=subprocess.PIPE, stderr=subprocess.PIPE, diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index c894774679..73c6759eec 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -29,11 +29,6 @@ _TEST_PORT = 13031 -# --------------------------------------------------------------------------- -# MockViewerPublisher -# --------------------------------------------------------------------------- - - class MockViewerPublisher: """Python mirror of the Rust WsPublisher in dimos-viewer. @@ -55,10 +50,6 @@ def __init__(self, url: str) -> None: self._ws: Any = None self._loop: asyncio.AbstractEventLoop | None = None - # ------------------------------------------------------------------ - # Context-manager interface - # ------------------------------------------------------------------ - def __enter__(self) -> "MockViewerPublisher": self._loop = asyncio.new_event_loop() self._ws = self._loop.run_until_complete(self._connect()) @@ -75,10 +66,6 @@ async def _connect(self) -> Any: return await ws_client.connect(self._url) - # ------------------------------------------------------------------ - # Send helpers (mirror of Rust WsPublisher methods) - # ------------------------------------------------------------------ - def send_click( self, x: float, @@ -138,11 +125,6 @@ def _send(self, msg: dict[str, Any]) -> None: self._loop.run_until_complete(self._ws.send(json.dumps(msg))) -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - def _collect(received: list[Any], done: threading.Event) -> Any: """Return a callback that appends to *received* and signals *done*.""" @@ -176,11 +158,6 @@ async def _probe() -> None: raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - class TestRerunWebSocketServerStartup: def test_server_binds_port(self) -> None: """After start(), the server must be reachable on the configured port.""" diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index ba0c953bd8..b374c739f0 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -77,10 +77,7 @@ def __init__(self, **kwargs: Any) -> None: self._ws_loop: asyncio.AbstractEventLoop | None = None self._server_thread: threading.Thread | None = None self._stop_event: asyncio.Event | None = None - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ + self._server_ready = threading.Event() @rpc def start(self) -> None: @@ -95,6 +92,9 @@ def start(self) -> None: @rpc def stop(self) -> None: + # Wait briefly for the server thread to initialise _stop_event so we + # don't silently skip the shutdown signal (race with _serve()). + self._server_ready.wait(timeout=5.0) if ( self._ws_loop is not None and not self._ws_loop.is_closed() @@ -103,10 +103,6 @@ def stop(self) -> None: self._ws_loop.call_soon_threadsafe(self._stop_event.set) super().stop() - # ------------------------------------------------------------------ - # Server - # ------------------------------------------------------------------ - def _run_server(self) -> None: """Entry point for the background server thread.""" self._ws_loop = asyncio.new_event_loop() @@ -120,6 +116,7 @@ async def _serve(self) -> None: import websockets.asyncio.server as ws_server self._stop_event = asyncio.Event() + self._server_ready.set() async with ws_server.serve( self._handle_client, @@ -140,10 +137,6 @@ async def _handle_client(self, websocket: Any) -> None: except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") - # ------------------------------------------------------------------ - # Message dispatch - # ------------------------------------------------------------------ - def _dispatch(self, raw: str | bytes) -> None: try: msg = json.loads(raw) @@ -151,6 +144,10 @@ def _dispatch(self, raw: str | bytes) -> None: logger.warning(f"RerunWebSocketServer: ignoring non-JSON message: {raw!r}") return + if not isinstance(msg, dict): + logger.warning(f"RerunWebSocketServer: expected JSON object, got {type(msg).__name__}") + return + msg_type = msg.get("type") if msg_type == "click": From ab4daea2529d2afe93c608fd0690b59dbdd673ef Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 17:13:58 -0700 Subject: [PATCH 17/53] consolidate viewer usage --- dimos/hardware/sensors/camera/module.py | 4 +- .../lidar/fastlio2/fastlio_blueprints.py | 39 ++++++++++++------- .../sensors/lidar/livox/livox_blueprints.py | 4 +- dimos/manipulation/blueprints.py | 4 +- dimos/manipulation/grasping/demo_grasping.py | 10 ++--- .../demo_object_scene_registration.py | 10 ++--- .../drone/blueprints/basic/drone_basic.py | 15 +------ .../blueprints/perceptive/unitree_g1_shm.py | 10 ++--- dimos/teleop/quest/blueprints.py | 4 +- 9 files changed, 47 insertions(+), 53 deletions(-) diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index b8165658d9..9c5623d141 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -32,7 +32,7 @@ from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier from dimos.spec import perception -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module def default_transform() -> Transform: @@ -120,5 +120,5 @@ def stop(self) -> None: demo_camera = autoconnect( CameraModule.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ) diff --git a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py index f3de842b46..b39dd7bcec 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py +++ b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py @@ -15,36 +15,45 @@ from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.lidar.fastlio2.module import FastLio2 from dimos.mapping.voxels import VoxelGridMapper -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module voxel_size = 0.05 mid360_fastlio = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=-1), - RerunBridgeModule.blueprint( - visual_override={ - "world/lidar": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), - } + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/lidar": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), + }, + }, ), ).global_config(n_workers=2, robot_model="mid360_fastlio2") mid360_fastlio_voxels = autoconnect( FastLio2.blueprint(), VoxelGridMapper.blueprint(publish_interval=1.0, voxel_size=voxel_size, carve_columns=False), - RerunBridgeModule.blueprint( - visual_override={ - "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), - "world/lidar": None, - } + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), + "world/lidar": None, + }, + }, ), ).global_config(n_workers=3, robot_model="mid360_fastlio2_voxels") mid360_fastlio_voxels_native = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=3.0), - RerunBridgeModule.blueprint( - visual_override={ - "world/lidar": None, - "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), - } + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/lidar": None, + "world/global_map": lambda grid: grid.to_rerun(voxel_size=voxel_size, mode="boxes"), + }, + }, ), ).global_config(n_workers=2, robot_model="mid360_fastlio2") diff --git a/dimos/hardware/sensors/lidar/livox/livox_blueprints.py b/dimos/hardware/sensors/lidar/livox/livox_blueprints.py index c8835b3e89..958af084e2 100644 --- a/dimos/hardware/sensors/lidar/livox/livox_blueprints.py +++ b/dimos/hardware/sensors/lidar/livox/livox_blueprints.py @@ -14,9 +14,9 @@ from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.lidar.livox.module import Mid360 -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module mid360 = autoconnect( Mid360.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ).global_config(n_workers=2, robot_model="mid360") diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index a9fb0fb44b..a2fd3389f0 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -46,7 +46,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.robot.foxglove_bridge import FoxgloveBridge # TODO: migrate to rerun +from dimos.visualization.vis_module import vis_module from dimos.utils.data import get_data @@ -409,7 +409,7 @@ def _make_piper_config( base_transform=_XARM_PERCEPTION_CAMERA_TRANSFORM, ), ObjectSceneRegistrationModule.blueprint(target_frame="world"), - FoxgloveBridge.blueprint(), # TODO: migrate to rerun + vis_module("foxglove"), ) .transports( { diff --git a/dimos/manipulation/grasping/demo_grasping.py b/dimos/manipulation/grasping/demo_grasping.py index 782283029b..f1ce67709e 100644 --- a/dimos/manipulation/grasping/demo_grasping.py +++ b/dimos/manipulation/grasping/demo_grasping.py @@ -14,15 +14,14 @@ # limitations under the License. from pathlib import Path -from dimos.agents.mcp.mcp_client import McpClient -from dimos.agents.mcp.mcp_server import McpServer +from dimos.agents.agent import Agent from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.manipulation.grasping.graspgen_module import graspgen from dimos.manipulation.grasping.grasping import GraspingModule from dimos.perception.detection.detectors.yoloe import YoloePromptMode from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.robot.foxglove_bridge import FoxgloveBridge +from dimos.visualization.vis_module import vis_module camera_module = RealSenseCamera.blueprint(enable_pointcloud=False) @@ -44,7 +43,6 @@ ("/tmp", "/tmp", "rw") ], # Grasp visualization debug standalone: python -m dimos.manipulation.grasping.visualize_grasps ), - FoxgloveBridge.blueprint(), - McpServer.blueprint(), - McpClient.blueprint(), + vis_module("foxglove"), + Agent.blueprint(), ).global_config(viewer="foxglove") diff --git a/dimos/perception/demo_object_scene_registration.py b/dimos/perception/demo_object_scene_registration.py index c6d8c96625..13fb26cbb5 100644 --- a/dimos/perception/demo_object_scene_registration.py +++ b/dimos/perception/demo_object_scene_registration.py @@ -13,14 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.agents.mcp.mcp_client import McpClient -from dimos.agents.mcp.mcp_server import McpServer +from dimos.agents.agent import Agent from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.hardware.sensors.camera.zed.compat import ZEDCamera from dimos.perception.detection.detectors.yoloe import YoloePromptMode from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.robot.foxglove_bridge import FoxgloveBridge +from dimos.visualization.vis_module import vis_module camera_choice = "zed" @@ -34,7 +33,6 @@ demo_object_scene_registration = autoconnect( camera_module, ObjectSceneRegistrationModule.blueprint(target_frame="world", prompt_mode=YoloePromptMode.LRPC), - FoxgloveBridge.blueprint(), - McpServer.blueprint(), - McpClient.blueprint(), + vis_module("foxglove"), + Agent.blueprint(), ).global_config(viewer="foxglove") diff --git a/dimos/robot/drone/blueprints/basic/drone_basic.py b/dimos/robot/drone/blueprints/basic/drone_basic.py index fbe6621ae1..c99c273cc2 100644 --- a/dimos/robot/drone/blueprints/basic/drone_basic.py +++ b/dimos/robot/drone/blueprints/basic/drone_basic.py @@ -20,9 +20,9 @@ from dimos.core.blueprints import autoconnect from dimos.core.global_config import global_config -from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.robot.drone.camera_module import DroneCameraModule from dimos.robot.drone.connection_module import DroneConnectionModule +from dimos.visualization.vis_module import vis_module from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule @@ -60,23 +60,12 @@ def _drone_rerun_blueprint() -> Any: _rerun_config = { "blueprint": _drone_rerun_blueprint, - "pubsubs": [LCM()], "static": { "world/tf/base_link": _static_drone_body, }, } -# Conditional visualization -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge - - _vis = FoxgloveBridge.blueprint() -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode - - _vis = RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **_rerun_config) -else: - _vis = autoconnect() +_vis = vis_module(global_config.viewer, rerun_config=_rerun_config) # Determine connection string based on replay flag connection_string = "udp:0.0.0.0:14550" diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py index 5b127fb697..9efe400895 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py @@ -17,10 +17,11 @@ from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE from dimos.core.blueprints import autoconnect +from dimos.core.global_config import global_config from dimos.core.transport import pSHMTransport from dimos.msgs.sensor_msgs.Image import Image -from dimos.robot.foxglove_bridge import FoxgloveBridge from dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1 import unitree_g1 +from dimos.visualization.vis_module import vis_module unitree_g1_shm = autoconnect( unitree_g1.transports( @@ -30,10 +31,9 @@ ), } ), - FoxgloveBridge.blueprint( - shm_channels=[ - "/color_image#sensor_msgs.Image", - ] + vis_module( + global_config.viewer, + foxglove_config={"shm_channels": ["/color_image#sensor_msgs.Image"]}, ), ) diff --git a/dimos/teleop/quest/blueprints.py b/dimos/teleop/quest/blueprints.py index d6367310de..1b67de3b75 100644 --- a/dimos/teleop/quest/blueprints.py +++ b/dimos/teleop/quest/blueprints.py @@ -26,12 +26,12 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.teleop.quest.quest_extensions import ArmTeleopModule from dimos.teleop.quest.quest_types import Buttons -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module # Arm teleop with press-and-hold engage (has rerun viz) teleop_quest_rerun = autoconnect( ArmTeleopModule.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ).transports( { ("left_controller_output", PoseStamped): LCMTransport("/teleop/left_delta", PoseStamped), From dc5f2f8265e0b3f9a71a2c94a2783d03e79414e3 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 17:20:16 -0700 Subject: [PATCH 18/53] consolidate WebsocketVisModule --- .../drone/blueprints/basic/drone_basic.py | 2 -- .../primitive/uintree_g1_primitive_no_nav.py | 3 --- .../go2/blueprints/basic/unitree_go2_basic.py | 2 -- .../go2/blueprints/basic/unitree_go2_fleet.py | 6 ++--- dimos/visualization/vis_module.py | 24 +++++++++++++++---- 5 files changed, 21 insertions(+), 16 deletions(-) diff --git a/dimos/robot/drone/blueprints/basic/drone_basic.py b/dimos/robot/drone/blueprints/basic/drone_basic.py index c99c273cc2..c60483cb0a 100644 --- a/dimos/robot/drone/blueprints/basic/drone_basic.py +++ b/dimos/robot/drone/blueprints/basic/drone_basic.py @@ -23,7 +23,6 @@ from dimos.robot.drone.camera_module import DroneCameraModule from dimos.robot.drone.connection_module import DroneConnectionModule from dimos.visualization.vis_module import vis_module -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule def _static_drone_body(rr: Any) -> list[Any]: @@ -81,7 +80,6 @@ def _drone_rerun_blueprint() -> Any: outdoor=False, ), DroneCameraModule.blueprint(camera_intrinsics=[1000.0, 1000.0, 960.0, 540.0]), - WebsocketVisModule.blueprint(), ) __all__ = [ diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index 2228dbfd66..220caff949 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -41,7 +41,6 @@ WavefrontFrontierExplorer, ) from dimos.visualization.vis_module import vis_module -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule def _convert_camera_info(camera_info: Any) -> Any: @@ -135,8 +134,6 @@ def _create_webcam() -> Webcam: VoxelGridMapper.blueprint(voxel_size=0.1), CostMapper.blueprint(), WavefrontFrontierExplorer.blueprint(), - # Visualization - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_g1") .transports( diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 282f813571..1e0f32d25c 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -25,7 +25,6 @@ from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.robot.unitree.go2.connection import GO2Connection from dimos.visualization.vis_module import vis_module -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule # Mac has some issue with high bandwidth UDP, so we use pSHMTransport for color_image # actually we can use pSHMTransport for all platforms, and for all streams @@ -116,7 +115,6 @@ def _go2_rerun_blueprint() -> Any: autoconnect( _with_vis, GO2Connection.blueprint(), - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py index 1c55f3e93c..0468cad40d 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_fleet.py @@ -22,15 +22,13 @@ from dimos.core.blueprints import autoconnect from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator -from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import with_vis +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import _with_vis from dimos.robot.unitree.go2.fleet_connection import Go2FleetConnection -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule unitree_go2_fleet = ( autoconnect( - with_vis, + _with_vis, Go2FleetConnection.blueprint(), - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index de786f67e8..688a6efb5b 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -30,8 +30,8 @@ def vis_module( """Create a visualization blueprint based on the selected viewer backend. Bundles the appropriate viewer module (Rerun or Foxglove) together with - the ``RerunWebSocketServer`` so that remote viewer connections (click, - teleop) work out of the box when using a Rerun backend. + the ``WebsocketVisModule`` and ``RerunWebSocketServer`` so that the web + dashboard and remote viewer connections work out of the box. Example usage:: @@ -48,6 +48,8 @@ def vis_module( }, ) """ + from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule + if foxglove_config is None: foxglove_config = {} if rerun_config is None: @@ -59,8 +61,11 @@ def vis_module( case "foxglove": from dimos.robot.foxglove_bridge import FoxgloveBridge - return autoconnect(FoxgloveBridge.blueprint(**foxglove_config)) - case "rerun" | "rerun-web" | "rerun-connect": + return autoconnect( + FoxgloveBridge.blueprint(**foxglove_config), + WebsocketVisModule.blueprint(), + ) + case "rerun" | "rerun-web": from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule from dimos.visualization.rerun.websocket_server import RerunWebSocketServer @@ -68,6 +73,15 @@ def vis_module( return autoconnect( RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), RerunWebSocketServer.blueprint(), + WebsocketVisModule.blueprint(), + ) + case "rerun-connect": + from dimos.visualization.rerun.bridge import RerunBridgeModule + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + + return autoconnect( + RerunBridgeModule.blueprint(viewer_mode="connect", **rerun_config), + RerunWebSocketServer.blueprint(), ) case _: - return autoconnect() + return autoconnect(WebsocketVisModule.blueprint()) From ba14725fefd2e5fdd8422e1e24c37b130b00dac5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 19:28:26 -0700 Subject: [PATCH 19/53] fix: address PR review - server ready race, path filter, skip guard - Move _server_ready.set() inside ws_server.serve() context so stop() waits for the port to actually bind before sending shutdown signal - Add /ws path filter to reject non-viewer WebSocket connections - Add pytest.mark.skipif for dimos-viewer binary test in CI - Fix import ordering in manipulation/blueprints.py --- dimos/manipulation/blueprints.py | 2 +- .../visualization/rerun/test_viewer_ws_e2e.py | 18 +++++++++++++++--- dimos/visualization/rerun/websocket_server.py | 5 ++++- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index a2fd3389f0..90e468aaf2 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -46,8 +46,8 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule -from dimos.visualization.vis_module import vis_module from dimos.utils.data import get_data +from dimos.visualization.vis_module import vis_module def _make_base_pose( diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py index 5275adb660..80c4743e61 100644 --- a/dimos/visualization/rerun/test_viewer_ws_e2e.py +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -30,11 +30,14 @@ import asyncio import json import os +import shutil import subprocess import threading import time from typing import Any +import pytest + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer _E2E_PORT = 13032 @@ -259,6 +262,13 @@ class TestViewerBinaryConnectMode: """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket client attempts to connect to our Python server.""" + @pytest.mark.skipif( + shutil.which("dimos-viewer") is None + or "--connect" not in subprocess.run( + ["dimos-viewer", "--help"], capture_output=True, text=True + ).stdout, + reason="dimos-viewer binary not installed or does not support --connect", + ) def test_viewer_ws_client_connects(self) -> None: """dimos-viewer --connect starts and its WS client connects to our server.""" server = _make_server() @@ -307,11 +317,13 @@ def _on_pt(pt: Any) -> None: except subprocess.TimeoutExpired: proc.kill() + stdout = proc.stdout.read().decode(errors="replace") if proc.stdout else "" stderr = proc.stderr.read().decode(errors="replace") if proc.stderr else "" server.stop() # The viewer should log that it is connecting to our WS URL. - # Even without a display, the log output appears before the GUI loop starts. - assert "ws://127.0.0.1" in stderr or proc.returncode is not None, ( - f"Viewer did not attempt WS connection. stderr:\n{stderr}" + # Check both stdout and stderr since log output destination varies. + combined = stdout + stderr + assert f"ws://127.0.0.1:{_E2E_PORT}" in combined, ( + f"Viewer did not attempt WS connection.\nstdout:\n{stdout}\nstderr:\n{stderr}" ) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index b374c739f0..16a292ca87 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -116,19 +116,22 @@ async def _serve(self) -> None: import websockets.asyncio.server as ws_server self._stop_event = asyncio.Event() - self._server_ready.set() async with ws_server.serve( self._handle_client, host=self.config.host, port=self.config.port, ): + self._server_ready.set() logger.info( f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" ) await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: + if hasattr(websocket, "request") and websocket.request.path != "/ws": + await websocket.close(1008, "Not Found") + return addr = websocket.remote_address logger.info(f"RerunWebSocketServer: viewer connected from {addr}") try: From f670f1272bb0b0e660c5fe9a56d7a78c8c169730 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 19:31:14 -0700 Subject: [PATCH 20/53] fix: set explicit ping interval/timeout on WebSocket server The default websockets ping_interval=20s + ping_timeout=20s was too aggressive. Increase both to 30s to give the viewer more time to respond, especially during brief network hiccups. --- dimos/visualization/rerun/websocket_server.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 16a292ca87..b12307c11a 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -121,6 +121,10 @@ async def _serve(self) -> None: self._handle_client, host=self.config.host, port=self.config.port, + # Ping every 30 s, allow 30 s for pong — generous enough to + # survive brief network hiccups while still detecting dead clients. + ping_interval=30, + ping_timeout=30, ): self._server_ready.set() logger.info( From 7545a5a164d11b99df7977ebff380d20d26eb91f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 23:10:36 -0700 Subject: [PATCH 21/53] fix(rerun-ws): log exception and unblock stop() on server startup failure If _serve() throws (e.g. port in use), _server_ready was never set, causing stop() to block for 5s. Now logs the exception and sets _server_ready in finally block. Revert: git revert HEAD --- dimos/visualization/rerun/websocket_server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index b12307c11a..e75df4eb25 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -109,7 +109,10 @@ def _run_server(self) -> None: asyncio.set_event_loop(self._ws_loop) try: self._ws_loop.run_until_complete(self._serve()) + except Exception: + logger.exception("RerunWebSocketServer failed to start") finally: + self._server_ready.set() # unblock stop() even on failure self._ws_loop.close() async def _serve(self) -> None: From 204d8b7938b5e203bcb3fe0779fcbec1690721ea Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 23:24:44 -0700 Subject: [PATCH 22/53] docs: add changes.md with fix descriptions and revert instructions --- changes.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 changes.md diff --git a/changes.md b/changes.md new file mode 100644 index 0000000000..d1e4b4b2e7 --- /dev/null +++ b/changes.md @@ -0,0 +1,19 @@ +# PR #1643 (rconnect) — Paul Review Fixes + +## Commits (local, not pushed) + +### 1. `81769d273` — Log exception + unblock stop() on startup failure +- If `_serve()` throws, `_server_ready` was never set → `stop()` blocked 5s +- Now logs exception and sets `_server_ready` in finally +- **Revert:** `git revert 81769d273` + +## Reviewer was wrong on +- `_server_ready` race — it IS set inside `async with` (after bind), not before +- `msg.get("x") or 0` — code already uses `msg.get("x", 0)` correctly + +## Not addressed (need Jeff's input) +- `vis_module` always bundling `RerunWebSocketServer` — opt-out design choice +- `LCM()` instantiated for non-rerun backends — wasted resource +- `rerun-connect` skipping `WebsocketVisModule` — intentional? +- Default `host = "0.0.0.0"` — intentional for remote viewer use case +- Hardcoded test ports — should use port=0 for parallel safety From 2c39685967bf6cc401c8b4437daa72a0d66b6a6d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 14:19:43 -0700 Subject: [PATCH 23/53] feat: vis_module helper + rerun bridge improvements - vis_module(): unified visualization module factory (rerun/foxglove/none) - Rerun bridge: connect mode serves gRPC, logs viewer connection hints - Rerun bridge: graceful fallback when native viewer unavailable - RerunWebSocketServer: WebSocket relay for dimos-viewer - camera/module.py: use vis_module instead of direct RerunBridgeModule - go2_basic: use vis_module pattern - utils/generic: add is_jetson() and get_local_ips() helpers --- dimos/hardware/sensors/camera/module.py | 4 +- .../go2/blueprints/basic/unitree_go2_basic.py | 25 +- dimos/utils/generic.py | 37 ++ dimos/visualization/rerun/bridge.py | 159 +++---- .../visualization/rerun/test_viewer_ws_e2e.py | 328 ++++++++++++++ .../rerun/test_websocket_server.py | 407 ++++++++++++++++++ dimos/visualization/rerun/websocket_server.py | 202 +++++++++ dimos/visualization/vis_module.py | 85 ++++ 8 files changed, 1131 insertions(+), 116 deletions(-) create mode 100644 dimos/visualization/rerun/test_viewer_ws_e2e.py create mode 100644 dimos/visualization/rerun/test_websocket_server.py create mode 100644 dimos/visualization/rerun/websocket_server.py create mode 100644 dimos/visualization/vis_module.py diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index b8165658d9..9c5623d141 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -32,7 +32,7 @@ from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier from dimos.spec import perception -from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.vis_module import vis_module def default_transform() -> Transform: @@ -120,5 +120,5 @@ def stop(self) -> None: demo_camera = autoconnect( CameraModule.blueprint(), - RerunBridgeModule.blueprint(), + vis_module("rerun"), ) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index f32561e11d..cae339e957 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -25,7 +25,6 @@ from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.robot.unitree.go2.connection import GO2Connection -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule # Mac has some issue with high bandwidth UDP, so we use pSHMTransport for color_image # actually we can use pSHMTransport for all platforms, and for all streams @@ -115,28 +114,20 @@ def _go2_rerun_blueprint() -> Any: } -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge +from dimos.visualization.vis_module import vis_module - with_vis = autoconnect( - _transports_base, - FoxgloveBridge.blueprint(shm_channels=["/color_image#sensor_msgs.Image"]), - ) -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode +_vis = vis_module( + viewer_backend=global_config.viewer, + rerun_config=rerun_config, + foxglove_config={"shm_channels": ["/color_image#sensor_msgs.Image"]}, +) - with_vis = autoconnect( - _transports_base, - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config), - ) -else: - with_vis = _transports_base +_with_vis = autoconnect(_transports_base, _vis) unitree_go2_basic = ( autoconnect( - with_vis, + _with_vis, GO2Connection.blueprint(), - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) diff --git a/dimos/utils/generic.py b/dimos/utils/generic.py index 84168ce057..3b8529089a 100644 --- a/dimos/utils/generic.py +++ b/dimos/utils/generic.py @@ -13,13 +13,50 @@ # limitations under the License. from collections.abc import Callable +import functools import hashlib import json import os +from pathlib import Path +import platform import string +import sys from typing import Any, Generic, TypeVar, overload import uuid + +@functools.lru_cache(maxsize=1) +def is_jetson() -> bool: + """Check if running on an NVIDIA Jetson device.""" + if sys.platform != "linux": + return False + # Check kernel release for Tegra (most lightweight) + if "tegra" in platform.release().lower(): + return True + # Check device tree (works in containers with proper mounts) + try: + return "nvidia,tegra" in Path("/proc/device-tree/compatible").read_text() + except (FileNotFoundError, PermissionError): + pass + # Check for L4T release file + return Path("/etc/nv_tegra_release").exists() + + +def get_local_ips() -> list[tuple[str, str]]: + """Return ``(ip, interface_name)`` for every non-loopback IPv4 address. + + Picks up physical, virtual, and VPN interfaces (including Tailscale). + """ + import psutil + + results: list[tuple[str, str]] = [] + for iface, addrs in psutil.net_if_addrs().items(): + for addr in addrs: + if addr.family.name == "AF_INET" and not addr.address.startswith("127."): + results.append((addr.address, iface)) + return results + + _T = TypeVar("_T") diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index de89c5d347..cb28840401 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -19,7 +19,6 @@ from collections.abc import Callable from dataclasses import field from functools import lru_cache -import subprocess import time from typing import ( Any, @@ -55,6 +54,43 @@ RERUN_GRPC_PORT = 9876 RERUN_WEB_PORT = 9090 +RERUN_WS_PORT = 3030 + + +def _log_viewer_connect_hints(connect_url: str) -> None: + """Log the dimos-viewer / rerun command users should run to connect.""" + import socket + + # Extract port from connect URL (e.g. "rerun+http://127.0.0.1:9877/proxy") + from dimos.utils.generic import get_local_ips + + local_ips = get_local_ips() + hostname = socket.gethostname() + + ws_url = f"ws://127.0.0.1:{RERUN_WS_PORT}/ws" + + lines = [ + "", + "=" * 60, + "Connect a Rerun viewer to this machine:", + "", + f" dimos-viewer --connect {connect_url} --ws-url {ws_url}", + "", + ] + if local_ips: + lines.append("From another machine on the network:") + for ip, iface in local_ips: + remote_connect = connect_url.replace("127.0.0.1", ip) + remote_ws = ws_url.replace("127.0.0.1", ip) + lines.append( + f" dimos-viewer --connect {remote_connect} --ws-url {remote_ws} # {iface}" + ) + lines.append("") + lines.append(f" hostname: {hostname}") + lines.append("=" * 60) + lines.append("") + + logger.info("\n".join(lines)) # TODO OUT visual annotations # @@ -130,34 +166,12 @@ def to_rerun(self) -> RerunData: ... ViewerMode = Literal["native", "web", "connect", "none"] -def _hex_to_rgba(hex_color: str) -> int: - """Convert '#RRGGBB' to a 0xRRGGBBAA int (fully opaque).""" - h = hex_color.lstrip("#") - return (int(h, 16) << 8) | 0xFF - - -def _with_graph_tab(bp: Blueprint) -> Blueprint: - """Add a Graph tab alongside the existing viewer layout without changing it.""" - import rerun.blueprint as rrb - - root = bp.root_container - return rrb.Blueprint( - rrb.Tabs( - root, - rrb.GraphView(origin="blueprint", name="Graph"), - ), - auto_layout=bp.auto_layout, - auto_views=bp.auto_views, - collapse_panels=bp.collapse_panels, - ) - - def _default_blueprint() -> Blueprint: """Default blueprint with black background and raised grid.""" import rerun as rr import rerun.blueprint as rrb - return rrb.Blueprint( + return rrb.Blueprint( # type: ignore[no-any-return] rrb.Spatial3DView( origin="world", background=rrb.Background(kind="SolidColor", color=[0, 0, 0]), @@ -224,10 +238,6 @@ class RerunBridgeModule(Module[Config]): default_config = Config - GV_SCALE = 100.0 # graphviz inches to rerun screen units - MODULE_RADIUS = 30.0 - CHANNEL_RADIUS = 20.0 - @lru_cache(maxsize=256) def _visual_override_for_entity_path( self, entity_path: str @@ -317,6 +327,7 @@ def start(self) -> None: rr.init("dimos") if self.config.viewer_mode == "native": + spawned = False try: import rerun_bindings @@ -325,6 +336,7 @@ def start(self) -> None: executable_name="dimos-viewer", memory_limit=self.config.memory_limit, ) + spawned = True except ImportError: pass # dimos-viewer not installed except Exception: @@ -332,16 +344,35 @@ def start(self) -> None: "dimos-viewer found but failed to spawn, falling back to stock rerun", exc_info=True, ) - rr.spawn(connect=True, memory_limit=self.config.memory_limit) + if not spawned: + try: + rr.spawn(connect=True, memory_limit=self.config.memory_limit) + except (RuntimeError, FileNotFoundError): + logger.warning( + "Rerun native viewer not available (headless?). " + "Bridge will continue without a viewer — data is still " + "accessible via rerun-connect or rerun-web.", + exc_info=True, + ) elif self.config.viewer_mode == "web": server_uri = rr.serve_grpc() rr.serve_web_viewer(connect_to=server_uri, open_browser=False) elif self.config.viewer_mode == "connect": - rr.connect_grpc(self.config.connect_url) + # Serve gRPC so external viewers (dimos-viewer) can connect to us. + # Extract the port from the connect_url to match what viewers expect. + from urllib.parse import urlparse + + parsed = urlparse(self.config.connect_url.replace("rerun+", "", 1)) + grpc_port = parsed.port or RERUN_GRPC_PORT + rr.serve_grpc( + grpc_port=grpc_port, + server_memory_limit=self.config.memory_limit, + ) + _log_viewer_connect_hints(self.config.connect_url) # "none" - just init, no viewer (connect externally) if self.config.blueprint: - rr.send_blueprint(_with_graph_tab(self.config.blueprint())) + rr.send_blueprint(self.config.blueprint()) # Start pubsubs and subscribe to all messages for pubsub in self.config.pubsubs: @@ -369,72 +400,6 @@ def _log_static(self) -> None: else: rr.log(entity_path, data, static=True) - @rpc - def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: - """Log a blueprint module graph from a Graphviz DOT string. - - Runs ``dot -Tplain`` to compute positions, then logs - ``rr.GraphNodes`` + ``rr.GraphEdges`` to the active recording. - - Args: - dot_code: The DOT-format graph (from ``introspection.blueprint.dot.render``). - module_names: List of module class names (to distinguish modules from channels). - """ - import rerun as rr - - try: - result = subprocess.run( - ["dot", "-Tplain"], input=dot_code, text=True, capture_output=True, timeout=30 - ) - except (FileNotFoundError, subprocess.TimeoutExpired): - return - if result.returncode != 0: - return - - node_ids: list[str] = [] - node_labels: list[str] = [] - node_colors: list[int] = [] - positions: list[tuple[float, float]] = [] - radii: list[float] = [] - edges: list[tuple[str, str]] = [] - module_set = set(module_names) - - for line in result.stdout.splitlines(): - if line.startswith("node "): - parts = line.split() - node_id = parts[1].strip('"') - x = float(parts[2]) * self.GV_SCALE - y = -float(parts[3]) * self.GV_SCALE - label = parts[6].strip('"') - color = parts[9].strip('"') - - node_ids.append(node_id) - node_labels.append(label) - positions.append((x, y)) - node_colors.append(_hex_to_rgba(color)) - radii.append(self.MODULE_RADIUS if node_id in module_set else self.CHANNEL_RADIUS) - - elif line.startswith("edge "): - parts = line.split() - edges.append((parts[1].strip('"'), parts[2].strip('"'))) - - if not node_ids: - return - - rr.log( - "blueprint", - rr.GraphNodes( - node_ids=node_ids, - labels=node_labels, - colors=node_colors, - positions=positions, - radii=radii, - show_labels=True, - ), - rr.GraphEdges(edges=edges, graph_type="directed"), - static=True, - ) - @rpc def stop(self) -> None: super().stop() diff --git a/dimos/visualization/rerun/test_viewer_ws_e2e.py b/dimos/visualization/rerun/test_viewer_ws_e2e.py new file mode 100644 index 0000000000..ea8351f2f6 --- /dev/null +++ b/dimos/visualization/rerun/test_viewer_ws_e2e.py @@ -0,0 +1,328 @@ +# Copyright 2025-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. + +"""End-to-end test: dimos-viewer (headless) → WebSocket → RerunWebSocketServer. + +dimos-viewer is started in ``--connect`` mode so it initialises its WebSocket +client. The viewer needs a gRPC proxy to connect to; we give it a non-existent +one so the viewer starts up anyway but produces no visualisation. The important +part is that the WebSocket client inside the viewer tries to connect to +``ws://127.0.0.1:/ws``. + +Because the viewer is a native GUI application it cannot run headlessly in CI +without a display. This test therefore verifies the connection at the protocol +level by using the ``RerunWebSocketServer`` module directly as the server and +injecting synthetic JSON messages that mimic what the viewer would send once a +user clicks in the 3D viewport. +""" + +import asyncio +import json +import os +import shutil +import subprocess +import threading +import time +from typing import Any + +import pytest + +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +_E2E_PORT = 13032 + + +def _make_server(port: int = _E2E_PORT) -> RerunWebSocketServer: + return RerunWebSocketServer(port=port) + + +def _wait_for_server(port: int, timeout: float = 5.0) -> None: + import websockets.asyncio.client as ws_client + + async def _probe() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): + pass + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + asyncio.run(_probe()) + return + except Exception: + time.sleep(0.05) + raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") + + +def _send_messages(port: int, messages: list[dict[str, Any]], *, delay: float = 0.05) -> None: + import websockets.asyncio.client as ws_client + + async def _run() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws") as ws: + for msg in messages: + await ws.send(json.dumps(msg)) + await asyncio.sleep(delay) + + asyncio.run(_run()) + + +class TestViewerProtocolE2E: + """Verify the full Python-server side of the viewer ↔ DimOS protocol. + + These tests use the ``RerunWebSocketServer`` as the server and a dummy + WebSocket client (playing the role of dimos-viewer) to inject messages. + They confirm every message type is correctly routed and that only click + messages produce stream publishes. + """ + + def test_viewer_click_reaches_stream(self) -> None: + """A viewer click message received over WebSocket publishes PointStamped.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + done.set() + + server.clicked_point.subscribe(_on_pt) + + _send_messages( + _E2E_PORT, + [ + { + "type": "click", + "x": 10.0, + "y": 20.0, + "z": 0.5, + "entity_path": "/world/robot", + "timestamp_ms": 42000, + } + ], + ) + + done.wait(timeout=3.0) + server.stop() + + assert len(received) == 1 + pt = received[0] + assert abs(pt.x - 10.0) < 1e-9 + assert abs(pt.y - 20.0) < 1e-9 + assert abs(pt.z - 0.5) < 1e-9 + assert pt.frame_id == "/world/robot" + assert abs(pt.ts - 42.0) < 1e-6 + + def test_viewer_keyboard_twist_no_publish(self) -> None: + """Twist messages from keyboard control do not publish clicked_point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + server.clicked_point.subscribe(received.append) + + _send_messages( + _E2E_PORT, + [ + { + "type": "twist", + "linear_x": 0.5, + "linear_y": 0.0, + "linear_z": 0.0, + "angular_x": 0.0, + "angular_y": 0.0, + "angular_z": 0.8, + } + ], + ) + + server.stop() + assert received == [] + + def test_viewer_stop_no_publish(self) -> None: + """Stop messages do not publish clicked_point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + server.clicked_point.subscribe(received.append) + + _send_messages(_E2E_PORT, [{"type": "stop"}]) + + server.stop() + assert received == [] + + def test_full_viewer_session_sequence(self) -> None: + """Realistic session: connect, heartbeats, click, WASD, stop → one point.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + done.set() + + server.clicked_point.subscribe(_on_pt) + + _send_messages( + _E2E_PORT, + [ + # Initial heartbeats (viewer connects and starts 1 Hz heartbeat) + {"type": "heartbeat", "timestamp_ms": 1000}, + {"type": "heartbeat", "timestamp_ms": 2000}, + # User clicks a point in the 3D viewport + { + "type": "click", + "x": 3.14, + "y": 2.71, + "z": 1.41, + "entity_path": "/world", + "timestamp_ms": 3000, + }, + # User presses W (forward) + { + "type": "twist", + "linear_x": 0.5, + "linear_y": 0.0, + "linear_z": 0.0, + "angular_x": 0.0, + "angular_y": 0.0, + "angular_z": 0.0, + }, + # User releases W + {"type": "stop"}, + # Another heartbeat + {"type": "heartbeat", "timestamp_ms": 4000}, + ], + delay=0.2, + ) + + done.wait(timeout=3.0) + server.stop() + + assert len(received) == 1, f"Expected exactly 1 click, got {len(received)}" + pt = received[0] + assert abs(pt.x - 3.14) < 1e-9 + assert abs(pt.y - 2.71) < 1e-9 + assert abs(pt.z - 1.41) < 1e-9 + + def test_reconnect_after_disconnect(self) -> None: + """Server keeps accepting new connections after a client disconnects.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + all_done = threading.Event() + + def _on_pt(pt: Any) -> None: + received.append(pt) + if len(received) >= 2: + all_done.set() + + server.clicked_point.subscribe(_on_pt) + + # First connection — send one click and disconnect + _send_messages( + _E2E_PORT, + [{"type": "click", "x": 1.0, "y": 0.0, "z": 0.0, "entity_path": "", "timestamp_ms": 0}], + ) + + # Second connection (simulating viewer reconnect) — send another click + _send_messages( + _E2E_PORT, + [{"type": "click", "x": 2.0, "y": 0.0, "z": 0.0, "entity_path": "", "timestamp_ms": 0}], + ) + + all_done.wait(timeout=5.0) + server.stop() + + xs = sorted(pt.x for pt in received) + assert xs == [1.0, 2.0], f"Unexpected xs: {xs}" + + +class TestViewerBinaryConnectMode: + """Smoke test: dimos-viewer binary starts in --connect mode and its WebSocket + client attempts to connect to our Python server.""" + + @pytest.mark.skipif( + shutil.which("dimos-viewer") is None + or "--connect" + not in subprocess.run(["dimos-viewer", "--help"], capture_output=True, text=True).stdout, + reason="dimos-viewer binary not installed or does not support --connect", + ) + def test_viewer_ws_client_connects(self) -> None: + """dimos-viewer --connect starts and its WS client connects to our server.""" + server = _make_server() + server.start() + _wait_for_server(_E2E_PORT) + + received: list[Any] = [] + + def _on_pt(pt: Any) -> None: + received.append(pt) + + server.clicked_point.subscribe(_on_pt) + + # Start dimos-viewer in --connect mode, pointing it at a non-existent gRPC + # proxy (it will fail to stream data, but that's fine) and at our WS server. + # Use DISPLAY="" to prevent it from opening a window (it will exit quickly + # without a display, but the WebSocket connection happens before the GUI loop). + proc = subprocess.Popen( + [ + "dimos-viewer", + "--connect", + f"--ws-url=ws://127.0.0.1:{_E2E_PORT}/ws", + ], + env={ + **os.environ, + "DISPLAY": "", + }, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + # Give the viewer up to 5 s to connect its WebSocket client to our server. + # We detect the connection by waiting for the server to accept a client. + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline: + # Check if any connection was established by sending a message and + # verifying the viewer is still running. + if proc.poll() is not None: + # Viewer exited (expected without a display) — check if it connected first. + break + time.sleep(0.1) + + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + + stdout = proc.stdout.read().decode(errors="replace") if proc.stdout else "" + stderr = proc.stderr.read().decode(errors="replace") if proc.stderr else "" + server.stop() + + # The viewer should log that it is connecting to our WS URL. + # Check both stdout and stderr since log output destination varies. + combined = stdout + stderr + assert f"ws://127.0.0.1:{_E2E_PORT}" in combined, ( + f"Viewer did not attempt WS connection.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py new file mode 100644 index 0000000000..73c6759eec --- /dev/null +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -0,0 +1,407 @@ +# Copyright 2025-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. + +"""Tests for RerunWebSocketServer. + +Uses ``MockViewerPublisher`` to simulate dimos-viewer sending events, matching +the exact JSON protocol used by the Rust ``WsPublisher`` in the viewer. +""" + +import asyncio +import json +import threading +import time +from typing import Any + +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +_TEST_PORT = 13031 + + +class MockViewerPublisher: + """Python mirror of the Rust WsPublisher in dimos-viewer. + + Connects to a running ``RerunWebSocketServer`` and exposes the same + ``send_click`` / ``send_twist`` / ``send_stop`` / ``send_heartbeat`` + API that the real viewer uses. Useful for unit tests that need to + exercise the server without a real viewer binary. + + Usage:: + + with MockViewerPublisher("ws://127.0.0.1:13031/ws") as pub: + pub.send_click(1.0, 2.0, 0.0, "/world", timestamp_ms=1000) + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.send_stop() + """ + + def __init__(self, url: str) -> None: + self._url = url + self._ws: Any = None + self._loop: asyncio.AbstractEventLoop | None = None + + def __enter__(self) -> "MockViewerPublisher": + self._loop = asyncio.new_event_loop() + self._ws = self._loop.run_until_complete(self._connect()) + return self + + def __exit__(self, *_: Any) -> None: + if self._ws is not None and self._loop is not None: + self._loop.run_until_complete(self._ws.close()) + if self._loop is not None: + self._loop.close() + + async def _connect(self) -> Any: + import websockets.asyncio.client as ws_client + + return await ws_client.connect(self._url) + + def send_click( + self, + x: float, + y: float, + z: float, + entity_path: str = "", + timestamp_ms: int = 0, + ) -> None: + """Send a click event — matches viewer SelectionChange handler output.""" + self._send( + { + "type": "click", + "x": x, + "y": y, + "z": z, + "entity_path": entity_path, + "timestamp_ms": timestamp_ms, + } + ) + + def send_twist( + self, + linear_x: float, + linear_y: float, + linear_z: float, + angular_x: float, + angular_y: float, + angular_z: float, + ) -> None: + """Send a twist (WASD keyboard) event.""" + self._send( + { + "type": "twist", + "linear_x": linear_x, + "linear_y": linear_y, + "linear_z": linear_z, + "angular_x": angular_x, + "angular_y": angular_y, + "angular_z": angular_z, + } + ) + + def send_stop(self) -> None: + """Send a stop event (Space bar or key release).""" + self._send({"type": "stop"}) + + def send_heartbeat(self, timestamp_ms: int = 0) -> None: + """Send a heartbeat (1 Hz keepalive from viewer).""" + self._send({"type": "heartbeat", "timestamp_ms": timestamp_ms}) + + def flush(self, delay: float = 0.1) -> None: + """Wait briefly so the server processes queued messages.""" + time.sleep(delay) + + def _send(self, msg: dict[str, Any]) -> None: + assert self._loop is not None and self._ws is not None, "Not connected" + self._loop.run_until_complete(self._ws.send(json.dumps(msg))) + + +def _collect(received: list[Any], done: threading.Event) -> Any: + """Return a callback that appends to *received* and signals *done*.""" + + def _cb(msg: Any) -> None: + received.append(msg) + done.set() + + return _cb + + +def _make_module(port: int = _TEST_PORT) -> RerunWebSocketServer: + return RerunWebSocketServer(port=port) + + +def _wait_for_server(port: int, timeout: float = 3.0) -> None: + """Block until the WebSocket server accepts an upgrade handshake.""" + + async def _probe() -> None: + import websockets.asyncio.client as ws_client + + async with ws_client.connect(f"ws://127.0.0.1:{port}/ws"): + pass + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + asyncio.run(_probe()) + return + except Exception: + time.sleep(0.05) + raise TimeoutError(f"Server on port {port} did not become ready within {timeout}s") + + +class TestRerunWebSocketServerStartup: + def test_server_binds_port(self) -> None: + """After start(), the server must be reachable on the configured port.""" + mod = _make_module() + mod.start() + try: + _wait_for_server(_TEST_PORT) + finally: + mod.stop() + + def test_stop_is_idempotent(self) -> None: + """Calling stop() twice must not raise.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + mod.stop() + mod.stop() + + +class TestClickMessages: + def test_click_publishes_point_stamped(self) -> None: + """A single click publishes one PointStamped with correct coords.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(1.5, 2.5, 0.0, "/world", timestamp_ms=1000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + pt = received[0] + assert abs(pt.x - 1.5) < 1e-9 + assert abs(pt.y - 2.5) < 1e-9 + assert abs(pt.z - 0.0) < 1e-9 + + def test_click_sets_frame_id_from_entity_path(self) -> None: + """entity_path is stored as frame_id on the published PointStamped.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(0.0, 0.0, 0.0, "/robot/base", timestamp_ms=2000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + assert received and received[0].frame_id == "/robot/base" + + def test_click_timestamp_converted_from_ms(self) -> None: + """timestamp_ms is converted to seconds on PointStamped.ts.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.clicked_point.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(0.0, 0.0, 0.0, "", timestamp_ms=5000) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + assert received and abs(received[0].ts - 5.0) < 1e-6 + + def test_multiple_clicks_all_published(self) -> None: + """A burst of clicks all arrive on the stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + all_arrived = threading.Event() + + def _cb(pt: Any) -> None: + received.append(pt) + if len(received) >= 3: + all_arrived.set() + + mod.clicked_point.subscribe(_cb) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_click(1.0, 0.0, 0.0) + pub.send_click(2.0, 0.0, 0.0) + pub.send_click(3.0, 0.0, 0.0) + pub.flush() + + all_arrived.wait(timeout=3.0) + mod.stop() + + assert sorted(pt.x for pt in received) == [1.0, 2.0, 3.0] + + +class TestNonClickMessages: + def test_heartbeat_does_not_publish(self) -> None: + """Heartbeat messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_heartbeat(9999) + pub.flush() + + mod.stop() + assert received == [] + + def test_twist_does_not_publish_clicked_point(self) -> None: + """Twist messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.flush() + + mod.stop() + assert received == [] + + def test_stop_does_not_publish_clicked_point(self) -> None: + """Stop messages must not trigger a clicked_point publish.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + mod.clicked_point.subscribe(received.append) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_stop() + pub.flush() + + mod.stop() + assert received == [] + + def test_twist_publishes_on_tele_cmd_vel(self) -> None: + """Twist messages publish a Twist on the tele_cmd_vel stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + tw = received[0] + assert abs(tw.linear.x - 0.5) < 1e-9 + assert abs(tw.angular.z - 0.8) < 1e-9 + + def test_stop_publishes_zero_twist_on_tele_cmd_vel(self) -> None: + """Stop messages publish a zero Twist on the tele_cmd_vel stream.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + received: list[Any] = [] + done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(received, done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_stop() + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + tw = received[0] + assert tw.is_zero() + + def test_invalid_json_does_not_crash(self) -> None: + """Malformed JSON is silently dropped; server stays alive.""" + import websockets.asyncio.client as ws_client + + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + async def _send_bad() -> None: + async with ws_client.connect(f"ws://127.0.0.1:{_TEST_PORT}/ws") as ws: + await ws.send("this is not json {{") + await asyncio.sleep(0.1) + await ws.send(json.dumps({"type": "heartbeat", "timestamp_ms": 0})) + await asyncio.sleep(0.1) + + asyncio.run(_send_bad()) + mod.stop() + + def test_mixed_message_sequence(self) -> None: + """Realistic sequence: heartbeat → click → twist → stop publishes one point.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + # Subscribe before sending so we don't race against the click dispatch. + received: list[Any] = [] + done = threading.Event() + + def _cb(pt: Any) -> None: + received.append(pt) + done.set() + + mod.clicked_point.subscribe(_cb) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_heartbeat(1000) + pub.send_click(7.0, 8.0, 9.0, "/map", timestamp_ms=1100) + pub.send_twist(0.3, 0.0, 0.0, 0.0, 0.0, 0.2) + pub.send_stop() + pub.flush() + + done.wait(timeout=2.0) + mod.stop() + + assert len(received) == 1 + assert abs(received[0].x - 7.0) < 1e-9 + assert abs(received[0].y - 8.0) < 1e-9 + assert abs(received[0].z - 9.0) < 1e-9 diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py new file mode 100644 index 0000000000..9431d1f00a --- /dev/null +++ b/dimos/visualization/rerun/websocket_server.py @@ -0,0 +1,202 @@ +# Copyright 2025-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. + +"""WebSocket server module that receives events from dimos-viewer. + +When dimos-viewer is started with ``--connect``, LCM multicast is unavailable +across machines. The viewer falls back to sending click, twist, and stop events +as JSON over a WebSocket connection. This module acts as the server-side +counterpart: it listens for those connections and translates incoming messages +into DimOS stream publishes. + +Message format (newline-delimited JSON, ``"type"`` discriminant): + + {"type":"heartbeat","timestamp_ms":1234567890} + {"type":"click","x":1.0,"y":2.0,"z":3.0,"entity_path":"/world","timestamp_ms":1234567890} + {"type":"twist","linear_x":0.5,"linear_y":0.0,"linear_z":0.0, + "angular_x":0.0,"angular_y":0.0,"angular_z":0.8} + {"type":"stop"} +""" + +import asyncio +import json +import threading +from typing import Any + +from dimos_lcm.std_msgs import Bool # type: ignore[import-untyped] +import websockets + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out +from dimos.msgs.geometry_msgs.PointStamped import PointStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class Config(ModuleConfig): + # Intentionally binds 0.0.0.0 by default so the viewer can connect from + # any machine on the network (the typical robot deployment scenario). + host: str = "0.0.0.0" + port: int = 3030 + + +class RerunWebSocketServer(Module[Config]): + """Receives dimos-viewer WebSocket events and publishes them as DimOS streams. + + The viewer connects to this module (not the other way around) when running + in ``--connect`` mode. Each click event is converted to a ``PointStamped`` + and published on the ``clicked_point`` stream so downstream modules (e.g. + ``ReplanningAStarPlanner``) can consume it without modification. + + Outputs: + clicked_point: 3-D world-space point from the most recent viewer click. + tele_cmd_vel: Twist velocity commands from keyboard teleop, including stop events. + """ + + default_config = Config + + clicked_point: Out[PointStamped] + tele_cmd_vel: Out[Twist] + stop_explore_cmd: Out[Bool] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._teleop_active = False + self._ws_loop: asyncio.AbstractEventLoop | None = None + self._server_thread: threading.Thread | None = None + self._stop_event: asyncio.Event | None = None + self._server_ready = threading.Event() + + @rpc + def start(self) -> None: + super().start() + self._server_thread = threading.Thread( + target=self._run_server, daemon=True, name="rerun-ws-server" + ) + self._server_thread.start() + logger.info( + f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws" + ) + + @rpc + def stop(self) -> None: + # Wait briefly for the server thread to initialise _stop_event so we + # don't silently skip the shutdown signal (race with _serve()). + self._server_ready.wait(timeout=5.0) + if ( + self._ws_loop is not None + and not self._ws_loop.is_closed() + and self._stop_event is not None + ): + self._ws_loop.call_soon_threadsafe(self._stop_event.set) + super().stop() + + def _run_server(self) -> None: + """Entry point for the background server thread.""" + self._ws_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._ws_loop) + try: + self._ws_loop.run_until_complete(self._serve()) + finally: + self._ws_loop.close() + + async def _serve(self) -> None: + import websockets.asyncio.server as ws_server + + self._stop_event = asyncio.Event() + + async with ws_server.serve( + self._handle_client, + host=self.config.host, + port=self.config.port, + # Ping every 30 s, allow 30 s for pong — generous enough to + # survive brief network hiccups while still detecting dead clients. + ping_interval=30, + ping_timeout=30, + ): + self._server_ready.set() + logger.info( + f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" + ) + await self._stop_event.wait() + + async def _handle_client(self, websocket: Any) -> None: + if hasattr(websocket, "request") and websocket.request.path != "/ws": + await websocket.close(1008, "Not Found") + return + addr = websocket.remote_address + logger.info(f"RerunWebSocketServer: viewer connected from {addr}") + try: + async for raw in websocket: + self._dispatch(raw) + except websockets.ConnectionClosed as exc: + logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") + + def _dispatch(self, raw: str | bytes) -> None: + try: + msg = json.loads(raw) + except json.JSONDecodeError: + logger.warning(f"RerunWebSocketServer: ignoring non-JSON message: {raw!r}") + return + + if not isinstance(msg, dict): + logger.warning(f"RerunWebSocketServer: expected JSON object, got {type(msg).__name__}") + return + + msg_type = msg.get("type") + + if msg_type == "click": + pt = PointStamped( + x=float(msg.get("x", 0)), + y=float(msg.get("y", 0)), + z=float(msg.get("z", 0)), + ts=float(msg.get("timestamp_ms", 0)) / 1000.0, + frame_id=str(msg.get("entity_path", "")), + ) + logger.debug(f"RerunWebSocketServer: click → {pt}") + self.clicked_point.publish(pt) + + elif msg_type == "twist": + twist = Twist( + linear=Vector3( + float(msg.get("linear_x", 0)), + float(msg.get("linear_y", 0)), + float(msg.get("linear_z", 0)), + ), + angular=Vector3( + float(msg.get("angular_x", 0)), + float(msg.get("angular_y", 0)), + float(msg.get("angular_z", 0)), + ), + ) + logger.debug(f"RerunWebSocketServer: twist → {twist}") + if not self._teleop_active: + self._teleop_active = True + self.stop_explore_cmd.publish(Bool(data=True)) + self.tele_cmd_vel.publish(twist) + + elif msg_type == "stop": + logger.debug("RerunWebSocketServer: stop") + self._teleop_active = False + self.tele_cmd_vel.publish(Twist.zero()) + + elif msg_type == "heartbeat": + logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") + + else: + logger.warning(f"RerunWebSocketServer: unknown message type {msg_type!r}") diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py new file mode 100644 index 0000000000..ccd425a540 --- /dev/null +++ b/dimos/visualization/vis_module.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# Copyright 2025-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. + +"""Shared visualization module factory for all robot blueprints.""" + +from typing import Any + +from dimos.core.blueprints import Blueprint, autoconnect +from dimos.core.global_config import ViewerBackend +from dimos.protocol.pubsub.impl.lcmpubsub import LCM + + +def vis_module( + viewer_backend: ViewerBackend, + rerun_config: dict[str, Any] | None = None, + foxglove_config: dict[str, Any] | None = None, +) -> Blueprint: + """Create a visualization blueprint based on the selected viewer backend. + + Bundles the appropriate viewer module (Rerun or Foxglove) together with + the ``RerunWebSocketServer`` so that the dimos-viewer keyboard/click + events work out of the box. + + Example usage:: + + + from dimos.core.global_config import global_config + viz = vis_module( + global_config.viewer, + rerun_config={ + "visual_override": { + "world/camera_info": lambda ci: ci.to_rerun(...), + }, + "static": { + "world/tf/base_link": lambda rr: [rr.Boxes3D(...)], + }, + }, + ) + """ + from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + + if foxglove_config is None: + foxglove_config = {} + if rerun_config is None: + rerun_config = {} + rerun_config = {**rerun_config} + rerun_config.setdefault("pubsubs", [LCM()]) + + match viewer_backend: + case "foxglove": + from dimos.robot.foxglove_bridge import FoxgloveBridge + + return autoconnect( + FoxgloveBridge.blueprint(**foxglove_config), + RerunWebSocketServer.blueprint(), + ) + case "rerun" | "rerun-web": + from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule + + viewer_mode = _BACKEND_TO_MODE.get(viewer_backend, "native") + return autoconnect( + RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), + RerunWebSocketServer.blueprint(), + ) + case "rerun-connect": + from dimos.visualization.rerun.bridge import RerunBridgeModule + + return autoconnect( + RerunBridgeModule.blueprint(viewer_mode="connect", **rerun_config), + RerunWebSocketServer.blueprint(), + ) + case _: + return autoconnect(RerunWebSocketServer.blueprint()) From 3ecae4bd85019826e676b708928f756e993731e4 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 15:01:52 -0700 Subject: [PATCH 24/53] cleanup --- .../blueprints/perceptive/unitree_g1_shm.py | 10 +++++----- .../primitive/uintree_g1_primitive_no_nav.py | 19 ++++++------------- dimos/visualization/rerun/bridge.py | 1 + 3 files changed, 12 insertions(+), 18 deletions(-) diff --git a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py index 5b127fb697..721487d717 100644 --- a/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py +++ b/dimos/robot/unitree/g1/blueprints/perceptive/unitree_g1_shm.py @@ -17,10 +17,11 @@ from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE from dimos.core.blueprints import autoconnect +from dimos.core.global_config import global_config from dimos.core.transport import pSHMTransport from dimos.msgs.sensor_msgs.Image import Image -from dimos.robot.foxglove_bridge import FoxgloveBridge from dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1 import unitree_g1 +from dimos.visualization.vis_module import vis_module unitree_g1_shm = autoconnect( unitree_g1.transports( @@ -30,10 +31,9 @@ ), } ), - FoxgloveBridge.blueprint( - shm_channels=[ - "/color_image#sensor_msgs.Image", - ] + vis_module( + viewer_backend=global_config.viewer, + foxglove_config={"shm_channels": ["/color_image#sensor_msgs.Image"]}, ), ) diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index ff59c9b8ef..fc9ddc58d6 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -41,7 +41,6 @@ WavefrontFrontierExplorer, ) from dimos.protocol.pubsub.impl.lcmpubsub import LCM -from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule def _convert_camera_info(camera_info: Any) -> Any: @@ -109,18 +108,14 @@ def _g1_rerun_blueprint() -> Any: }, } -if global_config.viewer == "foxglove": - from dimos.robot.foxglove_bridge import FoxgloveBridge +from dimos.visualization.vis_module import vis_module - _with_vis = autoconnect(FoxgloveBridge.blueprint()) -elif global_config.viewer.startswith("rerun"): - from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode +_vis = vis_module( + viewer_backend=global_config.viewer, + rerun_config=rerun_config, +) - _with_vis = autoconnect( - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config) - ) -else: - _with_vis = autoconnect() +_with_vis = autoconnect(_vis) def _create_webcam() -> Webcam: @@ -155,8 +150,6 @@ def _create_webcam() -> Webcam: VoxelGridMapper.blueprint(voxel_size=0.1), CostMapper.blueprint(), WavefrontFrontierExplorer.blueprint(), - # Visualization - WebsocketVisModule.blueprint(), ) .global_config(n_workers=4, robot_model="unitree_g1") .transports( diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index cb28840401..86a964ba04 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -92,6 +92,7 @@ def _log_viewer_connect_hints(connect_url: str) -> None: logger.info("\n".join(lines)) + # TODO OUT visual annotations # # In the future it would be nice if modules can annotate their individual OUTs with (general or rerun specific) From 77e25eb23978cec689a5609d50a9b28b1617ee2d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 15:08:35 -0700 Subject: [PATCH 25/53] refine --- dimos/hardware/sensors/camera/module.py | 3 +- .../rerun/test_websocket_server.py | 80 ++++++++++++++++--- dimos/visualization/rerun/websocket_server.py | 5 ++ dimos/visualization/vis_module.py | 17 ++-- 4 files changed, 84 insertions(+), 21 deletions(-) diff --git a/dimos/hardware/sensors/camera/module.py b/dimos/hardware/sensors/camera/module.py index 9c5623d141..de6ee2293c 100644 --- a/dimos/hardware/sensors/camera/module.py +++ b/dimos/hardware/sensors/camera/module.py @@ -22,6 +22,7 @@ from dimos.agents.annotation import skill from dimos.core.blueprints import autoconnect from dimos.core.core import rpc +from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import CameraHardware @@ -120,5 +121,5 @@ def stop(self) -> None: demo_camera = autoconnect( CameraModule.blueprint(), - vis_module("rerun"), + vis_module(global_config.viewer), ) diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index 73c6759eec..cec85fbb11 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -272,15 +272,21 @@ def test_heartbeat_does_not_publish(self) -> None: mod.start() _wait_for_server(_TEST_PORT) - received: list[Any] = [] - mod.clicked_point.subscribe(received.append) + clicks: list[Any] = [] + twists: list[Any] = [] + twist_done = threading.Event() + mod.clicked_point.subscribe(clicks.append) + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_heartbeat(9999) + # Send a canary twist so we know the server processed everything + pub.send_stop() pub.flush() + twist_done.wait(timeout=2.0) mod.stop() - assert received == [] + assert clicks == [] def test_twist_does_not_publish_clicked_point(self) -> None: """Twist messages must not trigger a clicked_point publish.""" @@ -288,15 +294,19 @@ def test_twist_does_not_publish_clicked_point(self) -> None: mod.start() _wait_for_server(_TEST_PORT) - received: list[Any] = [] - mod.clicked_point.subscribe(received.append) + clicks: list[Any] = [] + twists: list[Any] = [] + twist_done = threading.Event() + mod.clicked_point.subscribe(clicks.append) + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) pub.flush() + twist_done.wait(timeout=2.0) mod.stop() - assert received == [] + assert clicks == [] def test_stop_does_not_publish_clicked_point(self) -> None: """Stop messages must not trigger a clicked_point publish.""" @@ -304,15 +314,19 @@ def test_stop_does_not_publish_clicked_point(self) -> None: mod.start() _wait_for_server(_TEST_PORT) - received: list[Any] = [] - mod.clicked_point.subscribe(received.append) + clicks: list[Any] = [] + twists: list[Any] = [] + twist_done = threading.Event() + mod.clicked_point.subscribe(clicks.append) + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_stop() pub.flush() + twist_done.wait(timeout=2.0) mod.stop() - assert received == [] + assert clicks == [] def test_twist_publishes_on_tele_cmd_vel(self) -> None: """Twist messages publish a Twist on the tele_cmd_vel stream.""" @@ -357,6 +371,54 @@ def test_stop_publishes_zero_twist_on_tele_cmd_vel(self) -> None: tw = received[0] assert tw.is_zero() + def test_twist_publishes_stop_explore_cmd_on_first_twist(self) -> None: + """First twist publishes Bool(data=True) on stop_explore_cmd; stop resets.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + explore_cmds: list[Any] = [] + twists: list[Any] = [] + first_done = threading.Event() + mod.stop_explore_cmd.subscribe(_collect(explore_cmds, first_done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.0) + pub.flush() + + first_done.wait(timeout=2.0) + assert len(explore_cmds) == 1 + assert explore_cmds[0].data is True + + # Second twist within same connection should NOT publish another stop_explore_cmd + twist_done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) + + pub.send_twist(0.3, 0.0, 0.0, 0.0, 0.0, 0.0) + pub.flush() + + twist_done.wait(timeout=2.0) + assert len(explore_cmds) == 1 # still just the first one + + # After stop + new twist within same connection, stop_explore_cmd should fire again + second_done = threading.Event() + + def _on_second(msg: Any) -> None: + explore_cmds.append(msg) + if len(explore_cmds) >= 2: + second_done.set() + + mod.stop_explore_cmd.subscribe(_on_second) + + pub.send_stop() + pub.send_twist(0.1, 0.0, 0.0, 0.0, 0.0, 0.0) + pub.flush() + + second_done.wait(timeout=2.0) + + mod.stop() + assert len(explore_cmds) >= 2 + def test_invalid_json_does_not_crash(self) -> None: """Malformed JSON is silently dropped; server stays alive.""" import websockets.asyncio.client as ws_client diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 9431d1f00a..51fbff8fab 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -112,7 +112,10 @@ def _run_server(self) -> None: asyncio.set_event_loop(self._ws_loop) try: self._ws_loop.run_until_complete(self._serve()) + except Exception: + logger.error("RerunWebSocketServer failed to start", exc_info=True) finally: + self._server_ready.set() # unblock stop() even on failure self._ws_loop.close() async def _serve(self) -> None: @@ -146,6 +149,8 @@ async def _handle_client(self, websocket: Any) -> None: self._dispatch(raw) except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") + finally: + self._teleop_active = False def _dispatch(self, raw: str | bytes) -> None: try: diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index ccd425a540..aab461ae22 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -19,7 +19,6 @@ from dimos.core.blueprints import Blueprint, autoconnect from dimos.core.global_config import ViewerBackend -from dimos.protocol.pubsub.impl.lcmpubsub import LCM def vis_module( @@ -55,8 +54,6 @@ def vis_module( foxglove_config = {} if rerun_config is None: rerun_config = {} - rerun_config = {**rerun_config} - rerun_config.setdefault("pubsubs", [LCM()]) match viewer_backend: case "foxglove": @@ -66,20 +63,18 @@ def vis_module( FoxgloveBridge.blueprint(**foxglove_config), RerunWebSocketServer.blueprint(), ) - case "rerun" | "rerun-web": + case "rerun" | "rerun-web" | "rerun-connect": + from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule + rerun_config = {**rerun_config} + rerun_config.setdefault("pubsubs", [LCM()]) viewer_mode = _BACKEND_TO_MODE.get(viewer_backend, "native") return autoconnect( RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), RerunWebSocketServer.blueprint(), ) - case "rerun-connect": - from dimos.visualization.rerun.bridge import RerunBridgeModule - - return autoconnect( - RerunBridgeModule.blueprint(viewer_mode="connect", **rerun_config), - RerunWebSocketServer.blueprint(), - ) + case "none": + return autoconnect() case _: return autoconnect(RerunWebSocketServer.blueprint()) From 3835a39b36171262107037f5473ba46d62d49b30 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 15:39:25 -0700 Subject: [PATCH 26/53] cleanup --- .../primitive/uintree_g1_primitive_no_nav.py | 3 +-- .../go2/blueprints/basic/unitree_go2_basic.py | 3 +-- dimos/utils/generic.py | 4 ++- dimos/visualization/rerun/bridge.py | 4 ++- dimos/visualization/rerun/websocket_server.py | 25 +++++++++---------- dimos/visualization/vis_module.py | 7 +++++- 6 files changed, 26 insertions(+), 20 deletions(-) diff --git a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py index fc9ddc58d6..17a9389a7c 100644 --- a/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py +++ b/dimos/robot/unitree/g1/blueprints/primitive/uintree_g1_primitive_no_nav.py @@ -41,6 +41,7 @@ WavefrontFrontierExplorer, ) from dimos.protocol.pubsub.impl.lcmpubsub import LCM +from dimos.visualization.vis_module import vis_module def _convert_camera_info(camera_info: Any) -> Any: @@ -108,8 +109,6 @@ def _g1_rerun_blueprint() -> Any: }, } -from dimos.visualization.vis_module import vis_module - _vis = vis_module( viewer_backend=global_config.viewer, rerun_config=rerun_config, diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index cae339e957..052e220f1a 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -25,6 +25,7 @@ from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.service.system_configurator.clock_sync import ClockSyncConfigurator from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.visualization.vis_module import vis_module # Mac has some issue with high bandwidth UDP, so we use pSHMTransport for color_image # actually we can use pSHMTransport for all platforms, and for all streams @@ -114,8 +115,6 @@ def _go2_rerun_blueprint() -> Any: } -from dimos.visualization.vis_module import vis_module - _vis = vis_module( viewer_backend=global_config.viewer, rerun_config=rerun_config, diff --git a/dimos/utils/generic.py b/dimos/utils/generic.py index 3b8529089a..6aa1859659 100644 --- a/dimos/utils/generic.py +++ b/dimos/utils/generic.py @@ -47,12 +47,14 @@ def get_local_ips() -> list[tuple[str, str]]: Picks up physical, virtual, and VPN interfaces (including Tailscale). """ + import socket + import psutil results: list[tuple[str, str]] = [] for iface, addrs in psutil.net_if_addrs().items(): for addr in addrs: - if addr.family.name == "AF_INET" and not addr.address.startswith("127."): + if addr.family == socket.AF_INET and not addr.address.startswith("127."): results.append((addr.address, iface)) return results diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 86a964ba04..e3fdf967f4 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -238,6 +238,7 @@ class RerunBridgeModule(Module[Config]): """ default_config = Config + _last_log: dict[str, float] = {} @lru_cache(maxsize=256) def _visual_override_for_entity_path( @@ -321,7 +322,7 @@ def start(self) -> None: super().start() - self._last_log: dict[str, float] = {} + self._last_log: dict[str, float] = {} # reset on each start logger.info("Rerun bridge starting", viewer_mode=self.config.viewer_mode) # Initialize and spawn Rerun viewer @@ -403,6 +404,7 @@ def _log_static(self) -> None: @rpc def stop(self) -> None: + self._visual_override_for_entity_path.cache_clear() super().stop() diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 51fbff8fab..d868a2920b 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -53,6 +53,7 @@ class Config(ModuleConfig): # any machine on the network (the typical robot deployment scenario). host: str = "0.0.0.0" port: int = 3030 + start_timeout: float = 10.0 # seconds to wait for the server to bind class RerunWebSocketServer(Module[Config]): @@ -76,7 +77,7 @@ class RerunWebSocketServer(Module[Config]): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - self._teleop_active = False + self._teleop_clients: set[int] = set() # ids of clients currently in teleop self._ws_loop: asyncio.AbstractEventLoop | None = None self._server_thread: threading.Thread | None = None self._stop_event: asyncio.Event | None = None @@ -89,15 +90,16 @@ def start(self) -> None: target=self._run_server, daemon=True, name="rerun-ws-server" ) self._server_thread.start() + self._server_ready.wait(timeout=self.config.start_timeout) logger.info( - f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws" + f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" ) @rpc def stop(self) -> None: # Wait briefly for the server thread to initialise _stop_event so we # don't silently skip the shutdown signal (race with _serve()). - self._server_ready.wait(timeout=5.0) + self._server_ready.wait(timeout=self.config.start_timeout) if ( self._ws_loop is not None and not self._ws_loop.is_closed() @@ -109,7 +111,6 @@ def stop(self) -> None: def _run_server(self) -> None: """Entry point for the background server thread.""" self._ws_loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._ws_loop) try: self._ws_loop.run_until_complete(self._serve()) except Exception: @@ -133,9 +134,6 @@ async def _serve(self) -> None: ping_timeout=30, ): self._server_ready.set() - logger.info( - f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" - ) await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -143,16 +141,17 @@ async def _handle_client(self, websocket: Any) -> None: await websocket.close(1008, "Not Found") return addr = websocket.remote_address + client_id = id(websocket) logger.info(f"RerunWebSocketServer: viewer connected from {addr}") try: async for raw in websocket: - self._dispatch(raw) + self._dispatch(raw, client_id) except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") finally: - self._teleop_active = False + self._teleop_clients.discard(client_id) - def _dispatch(self, raw: str | bytes) -> None: + def _dispatch(self, raw: str | bytes, client_id: int) -> None: try: msg = json.loads(raw) except json.JSONDecodeError: @@ -190,14 +189,14 @@ def _dispatch(self, raw: str | bytes) -> None: ), ) logger.debug(f"RerunWebSocketServer: twist → {twist}") - if not self._teleop_active: - self._teleop_active = True + if not self._teleop_clients: self.stop_explore_cmd.publish(Bool(data=True)) + self._teleop_clients.add(client_id) self.tele_cmd_vel.publish(twist) elif msg_type == "stop": logger.debug("RerunWebSocketServer: stop") - self._teleop_active = False + self._teleop_clients.discard(client_id) self.tele_cmd_vel.publish(Twist.zero()) elif msg_type == "heartbeat": diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index aab461ae22..8c124883b8 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -59,6 +59,8 @@ def vis_module( case "foxglove": from dimos.robot.foxglove_bridge import FoxgloveBridge + # WS server is included even with Foxglove so dimos-viewer + # keyboard/click events still reach the robot. return autoconnect( FoxgloveBridge.blueprint(**foxglove_config), RerunWebSocketServer.blueprint(), @@ -77,4 +79,7 @@ def vis_module( case "none": return autoconnect() case _: - return autoconnect(RerunWebSocketServer.blueprint()) + raise ValueError( + f"Unknown viewer_backend {viewer_backend!r}. " + f"Expected one of: rerun, rerun-web, rerun-connect, foxglove, none" + ) From 1c77b5987353d5a094abf403d132a68b0f04756e Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 15:51:51 -0700 Subject: [PATCH 27/53] cleanup --- dimos/utils/generic.py | 39 +++++++++ dimos/visualization/rerun/bridge.py | 28 ++++++- .../rerun/test_websocket_server.py | 80 ++++++++++++++++--- dimos/visualization/rerun/websocket_server.py | 64 +++++++++++---- 4 files changed, 185 insertions(+), 26 deletions(-) diff --git a/dimos/utils/generic.py b/dimos/utils/generic.py index 84168ce057..6aa1859659 100644 --- a/dimos/utils/generic.py +++ b/dimos/utils/generic.py @@ -13,13 +13,52 @@ # limitations under the License. from collections.abc import Callable +import functools import hashlib import json import os +from pathlib import Path +import platform import string +import sys from typing import Any, Generic, TypeVar, overload import uuid + +@functools.lru_cache(maxsize=1) +def is_jetson() -> bool: + """Check if running on an NVIDIA Jetson device.""" + if sys.platform != "linux": + return False + # Check kernel release for Tegra (most lightweight) + if "tegra" in platform.release().lower(): + return True + # Check device tree (works in containers with proper mounts) + try: + return "nvidia,tegra" in Path("/proc/device-tree/compatible").read_text() + except (FileNotFoundError, PermissionError): + pass + # Check for L4T release file + return Path("/etc/nv_tegra_release").exists() + + +def get_local_ips() -> list[tuple[str, str]]: + """Return ``(ip, interface_name)`` for every non-loopback IPv4 address. + + Picks up physical, virtual, and VPN interfaces (including Tailscale). + """ + import socket + + import psutil + + results: list[tuple[str, str]] = [] + for iface, addrs in psutil.net_if_addrs().items(): + for addr in addrs: + if addr.family == socket.AF_INET and not addr.address.startswith("127."): + results.append((addr.address, iface)) + return results + + _T = TypeVar("_T") diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index de89c5d347..a23877b08e 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -56,6 +56,7 @@ RERUN_GRPC_PORT = 9876 RERUN_WEB_PORT = 9090 + # TODO OUT visual annotations # # In the future it would be nice if modules can annotate their individual OUTs with (general or rerun specific) @@ -223,6 +224,7 @@ class RerunBridgeModule(Module[Config]): """ default_config = Config + _last_log: dict[str, float] = {} GV_SCALE = 100.0 # graphviz inches to rerun screen units MODULE_RADIUS = 30.0 @@ -317,6 +319,7 @@ def start(self) -> None: rr.init("dimos") if self.config.viewer_mode == "native": + spawned = False try: import rerun_bindings @@ -325,6 +328,7 @@ def start(self) -> None: executable_name="dimos-viewer", memory_limit=self.config.memory_limit, ) + spawned = True except ImportError: pass # dimos-viewer not installed except Exception: @@ -332,12 +336,31 @@ def start(self) -> None: "dimos-viewer found but failed to spawn, falling back to stock rerun", exc_info=True, ) - rr.spawn(connect=True, memory_limit=self.config.memory_limit) + if not spawned: + try: + rr.spawn(connect=True, memory_limit=self.config.memory_limit) + except (RuntimeError, FileNotFoundError): + logger.warning( + "Rerun native viewer not available (headless?). " + "Bridge will continue without a viewer — data is still " + "accessible via rerun-connect or rerun-web.", + exc_info=True, + ) elif self.config.viewer_mode == "web": server_uri = rr.serve_grpc() rr.serve_web_viewer(connect_to=server_uri, open_browser=False) elif self.config.viewer_mode == "connect": - rr.connect_grpc(self.config.connect_url) + # Serve gRPC so external viewers (dimos-viewer) can connect to us. + # Extract the port from the connect_url to match what viewers expect. + from urllib.parse import urlparse + + parsed = urlparse(self.config.connect_url.replace("rerun+", "", 1)) + grpc_port = parsed.port or RERUN_GRPC_PORT + rr.serve_grpc( + grpc_port=grpc_port, + server_memory_limit=self.config.memory_limit, + ) + logger.info(f"Rerun gRPC server ready at {self.config.connect_url}") # "none" - just init, no viewer (connect externally) if self.config.blueprint: @@ -437,6 +460,7 @@ def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: @rpc def stop(self) -> None: + self._visual_override_for_entity_path.cache_clear() super().stop() diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index 73c6759eec..cec85fbb11 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -272,15 +272,21 @@ def test_heartbeat_does_not_publish(self) -> None: mod.start() _wait_for_server(_TEST_PORT) - received: list[Any] = [] - mod.clicked_point.subscribe(received.append) + clicks: list[Any] = [] + twists: list[Any] = [] + twist_done = threading.Event() + mod.clicked_point.subscribe(clicks.append) + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_heartbeat(9999) + # Send a canary twist so we know the server processed everything + pub.send_stop() pub.flush() + twist_done.wait(timeout=2.0) mod.stop() - assert received == [] + assert clicks == [] def test_twist_does_not_publish_clicked_point(self) -> None: """Twist messages must not trigger a clicked_point publish.""" @@ -288,15 +294,19 @@ def test_twist_does_not_publish_clicked_point(self) -> None: mod.start() _wait_for_server(_TEST_PORT) - received: list[Any] = [] - mod.clicked_point.subscribe(received.append) + clicks: list[Any] = [] + twists: list[Any] = [] + twist_done = threading.Event() + mod.clicked_point.subscribe(clicks.append) + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.8) pub.flush() + twist_done.wait(timeout=2.0) mod.stop() - assert received == [] + assert clicks == [] def test_stop_does_not_publish_clicked_point(self) -> None: """Stop messages must not trigger a clicked_point publish.""" @@ -304,15 +314,19 @@ def test_stop_does_not_publish_clicked_point(self) -> None: mod.start() _wait_for_server(_TEST_PORT) - received: list[Any] = [] - mod.clicked_point.subscribe(received.append) + clicks: list[Any] = [] + twists: list[Any] = [] + twist_done = threading.Event() + mod.clicked_point.subscribe(clicks.append) + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_stop() pub.flush() + twist_done.wait(timeout=2.0) mod.stop() - assert received == [] + assert clicks == [] def test_twist_publishes_on_tele_cmd_vel(self) -> None: """Twist messages publish a Twist on the tele_cmd_vel stream.""" @@ -357,6 +371,54 @@ def test_stop_publishes_zero_twist_on_tele_cmd_vel(self) -> None: tw = received[0] assert tw.is_zero() + def test_twist_publishes_stop_explore_cmd_on_first_twist(self) -> None: + """First twist publishes Bool(data=True) on stop_explore_cmd; stop resets.""" + mod = _make_module() + mod.start() + _wait_for_server(_TEST_PORT) + + explore_cmds: list[Any] = [] + twists: list[Any] = [] + first_done = threading.Event() + mod.stop_explore_cmd.subscribe(_collect(explore_cmds, first_done)) + + with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: + pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.0) + pub.flush() + + first_done.wait(timeout=2.0) + assert len(explore_cmds) == 1 + assert explore_cmds[0].data is True + + # Second twist within same connection should NOT publish another stop_explore_cmd + twist_done = threading.Event() + mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) + + pub.send_twist(0.3, 0.0, 0.0, 0.0, 0.0, 0.0) + pub.flush() + + twist_done.wait(timeout=2.0) + assert len(explore_cmds) == 1 # still just the first one + + # After stop + new twist within same connection, stop_explore_cmd should fire again + second_done = threading.Event() + + def _on_second(msg: Any) -> None: + explore_cmds.append(msg) + if len(explore_cmds) >= 2: + second_done.set() + + mod.stop_explore_cmd.subscribe(_on_second) + + pub.send_stop() + pub.send_twist(0.1, 0.0, 0.0, 0.0, 0.0, 0.0) + pub.flush() + + second_done.wait(timeout=2.0) + + mod.stop() + assert len(explore_cmds) >= 2 + def test_invalid_json_does_not_crash(self) -> None: """Malformed JSON is silently dropped; server stays alive.""" import websockets.asyncio.client as ws_client diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index e75df4eb25..af19f4c100 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -34,6 +34,7 @@ import threading from typing import Any +from dimos_lcm.std_msgs import Bool # type: ignore[import-untyped] import websockets from dimos.core.core import rpc @@ -52,8 +53,13 @@ class Config(ModuleConfig): # any machine on the network (the typical robot deployment scenario). host: str = "0.0.0.0" port: int = 3030 + start_timeout: float = 10.0 # seconds to wait for the server to bind - +# QUALITY LEVEL: temporary +# ideally this would be part of the rerun bridge +# SUPER ideally this module shouldn't exist at all (we should just patch rerun properly) +# but for now, I just need this to get the g1 stuff working +# the vis_module manages when to add the RerunWebSocketServer as a module alongside the RerunBridgeModule class RerunWebSocketServer(Module[Config]): """Receives dimos-viewer WebSocket events and publishes them as DimOS streams. @@ -71,9 +77,11 @@ class RerunWebSocketServer(Module[Config]): clicked_point: Out[PointStamped] tele_cmd_vel: Out[Twist] + stop_explore_cmd: Out[Bool] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) + self._teleop_clients: set[int] = set() # ids of clients currently in teleop self._ws_loop: asyncio.AbstractEventLoop | None = None self._server_thread: threading.Thread | None = None self._stop_event: asyncio.Event | None = None @@ -86,15 +94,14 @@ def start(self) -> None: target=self._run_server, daemon=True, name="rerun-ws-server" ) self._server_thread.start() - logger.info( - f"RerunWebSocketServer starting on ws://{self.config.host}:{self.config.port}/ws" - ) + self._server_ready.wait(timeout=self.config.start_timeout) + self._log_connect_hints() @rpc def stop(self) -> None: # Wait briefly for the server thread to initialise _stop_event so we # don't silently skip the shutdown signal (race with _serve()). - self._server_ready.wait(timeout=5.0) + self._server_ready.wait(timeout=self.config.start_timeout) if ( self._ws_loop is not None and not self._ws_loop.is_closed() @@ -103,14 +110,40 @@ def stop(self) -> None: self._ws_loop.call_soon_threadsafe(self._stop_event.set) super().stop() + def _log_connect_hints(self) -> None: + """Log the WebSocket URL(s) that viewers should connect to.""" + import socket + + from dimos.utils.generic import get_local_ips + + local_ips = get_local_ips() + hostname = socket.gethostname() + ws_url = f"ws://127.0.0.1:{self.config.port}/ws" + + lines = [ + "", + "=" * 60, + f"RerunWebSocketServer listening on {ws_url}", + "", + ] + if local_ips: + lines.append("From another machine on the network:") + for ip, iface in local_ips: + lines.append(f" ws://{ip}:{self.config.port}/ws # {iface}") + lines.append("") + lines.append(f" hostname: {hostname}") + lines.append("=" * 60) + lines.append("") + + logger.info("\n".join(lines)) + def _run_server(self) -> None: """Entry point for the background server thread.""" self._ws_loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._ws_loop) try: self._ws_loop.run_until_complete(self._serve()) except Exception: - logger.exception("RerunWebSocketServer failed to start") + logger.error("RerunWebSocketServer failed to start", exc_info=True) finally: self._server_ready.set() # unblock stop() even on failure self._ws_loop.close() @@ -130,9 +163,6 @@ async def _serve(self) -> None: ping_timeout=30, ): self._server_ready.set() - logger.info( - f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" - ) await self._stop_event.wait() async def _handle_client(self, websocket: Any) -> None: @@ -140,14 +170,17 @@ async def _handle_client(self, websocket: Any) -> None: await websocket.close(1008, "Not Found") return addr = websocket.remote_address + client_id = id(websocket) logger.info(f"RerunWebSocketServer: viewer connected from {addr}") try: async for raw in websocket: - self._dispatch(raw) + self._dispatch(raw, client_id) except websockets.ConnectionClosed as exc: logger.debug(f"RerunWebSocketServer: client {addr} disconnected ({exc})") + finally: + self._teleop_clients.discard(client_id) - def _dispatch(self, raw: str | bytes) -> None: + def _dispatch(self, raw: str | bytes, client_id: int) -> None: try: msg = json.loads(raw) except json.JSONDecodeError: @@ -185,10 +218,14 @@ def _dispatch(self, raw: str | bytes) -> None: ), ) logger.debug(f"RerunWebSocketServer: twist → {twist}") + if not self._teleop_clients: + self.stop_explore_cmd.publish(Bool(data=True)) + self._teleop_clients.add(client_id) self.tele_cmd_vel.publish(twist) elif msg_type == "stop": logger.debug("RerunWebSocketServer: stop") + self._teleop_clients.discard(client_id) self.tele_cmd_vel.publish(Twist.zero()) elif msg_type == "heartbeat": @@ -196,6 +233,3 @@ def _dispatch(self, raw: str | bytes) -> None: else: logger.warning(f"RerunWebSocketServer: unknown message type {msg_type!r}") - - -rerun_ws_server = RerunWebSocketServer.blueprint From b1dcf0f953f91a4df98fc120d584a424b33f10fd Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 15:52:56 -0700 Subject: [PATCH 28/53] misc --- dimos/visualization/rerun/bridge.py | 39 +------------------ dimos/visualization/rerun/websocket_server.py | 31 +++++++++++++-- 2 files changed, 29 insertions(+), 41 deletions(-) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index e3fdf967f4..823735d0c5 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -54,43 +54,6 @@ RERUN_GRPC_PORT = 9876 RERUN_WEB_PORT = 9090 -RERUN_WS_PORT = 3030 - - -def _log_viewer_connect_hints(connect_url: str) -> None: - """Log the dimos-viewer / rerun command users should run to connect.""" - import socket - - # Extract port from connect URL (e.g. "rerun+http://127.0.0.1:9877/proxy") - from dimos.utils.generic import get_local_ips - - local_ips = get_local_ips() - hostname = socket.gethostname() - - ws_url = f"ws://127.0.0.1:{RERUN_WS_PORT}/ws" - - lines = [ - "", - "=" * 60, - "Connect a Rerun viewer to this machine:", - "", - f" dimos-viewer --connect {connect_url} --ws-url {ws_url}", - "", - ] - if local_ips: - lines.append("From another machine on the network:") - for ip, iface in local_ips: - remote_connect = connect_url.replace("127.0.0.1", ip) - remote_ws = ws_url.replace("127.0.0.1", ip) - lines.append( - f" dimos-viewer --connect {remote_connect} --ws-url {remote_ws} # {iface}" - ) - lines.append("") - lines.append(f" hostname: {hostname}") - lines.append("=" * 60) - lines.append("") - - logger.info("\n".join(lines)) # TODO OUT visual annotations @@ -370,7 +333,7 @@ def start(self) -> None: grpc_port=grpc_port, server_memory_limit=self.config.memory_limit, ) - _log_viewer_connect_hints(self.config.connect_url) + logger.info(f"Rerun gRPC server ready at {self.config.connect_url}") # "none" - just init, no viewer (connect externally) if self.config.blueprint: diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index d868a2920b..9275018e32 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -91,9 +91,7 @@ def start(self) -> None: ) self._server_thread.start() self._server_ready.wait(timeout=self.config.start_timeout) - logger.info( - f"RerunWebSocketServer listening on ws://{self.config.host}:{self.config.port}/ws" - ) + self._log_connect_hints() @rpc def stop(self) -> None: @@ -108,6 +106,33 @@ def stop(self) -> None: self._ws_loop.call_soon_threadsafe(self._stop_event.set) super().stop() + def _log_connect_hints(self) -> None: + """Log the WebSocket URL(s) that viewers should connect to.""" + import socket + + from dimos.utils.generic import get_local_ips + + local_ips = get_local_ips() + hostname = socket.gethostname() + ws_url = f"ws://127.0.0.1:{self.config.port}/ws" + + lines = [ + "", + "=" * 60, + f"RerunWebSocketServer listening on {ws_url}", + "", + ] + if local_ips: + lines.append("From another machine on the network:") + for ip, iface in local_ips: + lines.append(f" ws://{ip}:{self.config.port}/ws # {iface}") + lines.append("") + lines.append(f" hostname: {hostname}") + lines.append("=" * 60) + lines.append("") + + logger.info("\n".join(lines)) + def _run_server(self) -> None: """Entry point for the background server thread.""" self._ws_loop = asyncio.new_event_loop() From 26fa2fb8f6a675032d2ae2946066bf67a22f92ea Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 30 Mar 2026 16:39:23 -0700 Subject: [PATCH 29/53] restoregraph --- dimos/visualization/rerun/bridge.py | 98 ++++++++++++++++++++++++++++- 1 file changed, 97 insertions(+), 1 deletion(-) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 823735d0c5..843ae421f4 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -19,6 +19,7 @@ from collections.abc import Callable from dataclasses import field from functools import lru_cache +import subprocess import time from typing import ( Any, @@ -130,6 +131,30 @@ def to_rerun(self) -> RerunData: ... ViewerMode = Literal["native", "web", "connect", "none"] +def _hex_to_rgba(hex_color: str) -> int: + """Convert '#RRGGBB' to a 0xRRGGBBAA int (fully opaque).""" + h = hex_color.lstrip("#") + if len(h) == 6: + return int(h + "ff", 16) + return int(h[:8], 16) + + +def _with_graph_tab(bp: Blueprint) -> Blueprint: + """Add a Graph tab alongside the existing viewer layout without changing it.""" + import rerun.blueprint as rrb + + root = bp.root_container + return rrb.Blueprint( + rrb.Tabs( + root, + rrb.GraphView(origin="blueprint", name="Graph"), + ), + auto_layout=bp.auto_layout, + auto_views=bp.auto_views, + collapse_panels=bp.collapse_panels, + ) + + def _default_blueprint() -> Blueprint: """Default blueprint with black background and raised grid.""" import rerun as rr @@ -203,6 +228,11 @@ class RerunBridgeModule(Module[Config]): default_config = Config _last_log: dict[str, float] = {} + # Graphviz layout scale and node radii for blueprint graph + GV_SCALE = 100.0 + MODULE_RADIUS = 20.0 + CHANNEL_RADIUS = 12.0 + @lru_cache(maxsize=256) def _visual_override_for_entity_path( self, entity_path: str @@ -337,7 +367,7 @@ def start(self) -> None: # "none" - just init, no viewer (connect externally) if self.config.blueprint: - rr.send_blueprint(self.config.blueprint()) + rr.send_blueprint(_with_graph_tab(self.config.blueprint())) # Start pubsubs and subscribe to all messages for pubsub in self.config.pubsubs: @@ -365,6 +395,72 @@ def _log_static(self) -> None: else: rr.log(entity_path, data, static=True) + @rpc + def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: + """Log a blueprint module graph from a Graphviz DOT string. + + Runs ``dot -Tplain`` to compute positions, then logs + ``rr.GraphNodes`` + ``rr.GraphEdges`` to the active recording. + + Args: + dot_code: The DOT-format graph (from ``introspection.blueprint.dot.render``). + module_names: List of module class names (to distinguish modules from channels). + """ + import rerun as rr + + try: + result = subprocess.run( + ["dot", "-Tplain"], input=dot_code, text=True, capture_output=True, timeout=30 + ) + except (FileNotFoundError, subprocess.TimeoutExpired): + return + if result.returncode != 0: + return + + node_ids: list[str] = [] + node_labels: list[str] = [] + node_colors: list[int] = [] + positions: list[tuple[float, float]] = [] + radii: list[float] = [] + edges: list[tuple[str, str]] = [] + module_set = set(module_names) + + for line in result.stdout.splitlines(): + if line.startswith("node "): + parts = line.split() + node_id = parts[1].strip('"') + x = float(parts[2]) * self.GV_SCALE + y = -float(parts[3]) * self.GV_SCALE + label = parts[6].strip('"') + color = parts[9].strip('"') + + node_ids.append(node_id) + node_labels.append(label) + positions.append((x, y)) + node_colors.append(_hex_to_rgba(color)) + radii.append(self.MODULE_RADIUS if node_id in module_set else self.CHANNEL_RADIUS) + + elif line.startswith("edge "): + parts = line.split() + edges.append((parts[1].strip('"'), parts[2].strip('"'))) + + if not node_ids: + return + + rr.log( + "blueprint", + rr.GraphNodes( + node_ids=node_ids, + labels=node_labels, + colors=node_colors, + positions=positions, + radii=radii, + show_labels=True, + ), + rr.GraphEdges(edges=edges, graph_type="directed"), + static=True, + ) + @rpc def stop(self) -> None: self._visual_override_for_entity_path.cache_clear() From 2a81168f9a98f7aa95f1526e3b3387a77997de8d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Apr 2026 17:15:27 -0700 Subject: [PATCH 30/53] change name to be more generic --- dimos/visualization/rerun/websocket_server.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 9275018e32..e7319d176a 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -67,13 +67,14 @@ class RerunWebSocketServer(Module[Config]): Outputs: clicked_point: 3-D world-space point from the most recent viewer click. tele_cmd_vel: Twist velocity commands from keyboard teleop, including stop events. + stop_movement: Published when teleop starts — signals nav to cancel the active goal. """ default_config = Config clicked_point: Out[PointStamped] tele_cmd_vel: Out[Twist] - stop_explore_cmd: Out[Bool] + stop_movement: Out[Bool] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -215,7 +216,7 @@ def _dispatch(self, raw: str | bytes, client_id: int) -> None: ) logger.debug(f"RerunWebSocketServer: twist → {twist}") if not self._teleop_clients: - self.stop_explore_cmd.publish(Bool(data=True)) + self.stop_movement.publish(Bool(data=True)) self._teleop_clients.add(client_id) self.tele_cmd_vel.publish(twist) From fb926c0cfd362b9b0bb87cbf84724df39aac7636 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Apr 2026 18:34:39 -0700 Subject: [PATCH 31/53] del --- changes.md | 19 ------------------- 1 file changed, 19 deletions(-) delete mode 100644 changes.md diff --git a/changes.md b/changes.md deleted file mode 100644 index d1e4b4b2e7..0000000000 --- a/changes.md +++ /dev/null @@ -1,19 +0,0 @@ -# PR #1643 (rconnect) — Paul Review Fixes - -## Commits (local, not pushed) - -### 1. `81769d273` — Log exception + unblock stop() on startup failure -- If `_serve()` throws, `_server_ready` was never set → `stop()` blocked 5s -- Now logs exception and sets `_server_ready` in finally -- **Revert:** `git revert 81769d273` - -## Reviewer was wrong on -- `_server_ready` race — it IS set inside `async with` (after bind), not before -- `msg.get("x") or 0` — code already uses `msg.get("x", 0)` correctly - -## Not addressed (need Jeff's input) -- `vis_module` always bundling `RerunWebSocketServer` — opt-out design choice -- `LCM()` instantiated for non-rerun backends — wasted resource -- `rerun-connect` skipping `WebsocketVisModule` — intentional? -- Default `host = "0.0.0.0"` — intentional for remote viewer use case -- Hardcoded test ports — should use port=0 for parallel safety From 8e227ec43f3bbe6af9e806725581455f9e242a39 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Apr 2026 19:08:29 -0700 Subject: [PATCH 32/53] fix old Agent.blueprint()'s --- dimos/manipulation/grasping/demo_grasping.py | 6 ++++-- dimos/perception/demo_object_scene_registration.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/dimos/manipulation/grasping/demo_grasping.py b/dimos/manipulation/grasping/demo_grasping.py index f1ce67709e..9ccb0e69bc 100644 --- a/dimos/manipulation/grasping/demo_grasping.py +++ b/dimos/manipulation/grasping/demo_grasping.py @@ -14,7 +14,8 @@ # limitations under the License. from pathlib import Path -from dimos.agents.agent import Agent +from dimos.agents.mcp.mcp_client import McpClient +from dimos.agents.mcp.mcp_server import McpServer from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.manipulation.grasping.graspgen_module import graspgen @@ -44,5 +45,6 @@ ], # Grasp visualization debug standalone: python -m dimos.manipulation.grasping.visualize_grasps ), vis_module("foxglove"), - Agent.blueprint(), + McpServer.blueprint(), + McpClient.blueprint(), ).global_config(viewer="foxglove") diff --git a/dimos/perception/demo_object_scene_registration.py b/dimos/perception/demo_object_scene_registration.py index 13fb26cbb5..531093f619 100644 --- a/dimos/perception/demo_object_scene_registration.py +++ b/dimos/perception/demo_object_scene_registration.py @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.agents.agent import Agent +from dimos.agents.mcp.mcp_client import McpClient +from dimos.agents.mcp.mcp_server import McpServer from dimos.core.blueprints import autoconnect from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.hardware.sensors.camera.zed.compat import ZEDCamera @@ -34,5 +35,6 @@ camera_module, ObjectSceneRegistrationModule.blueprint(target_frame="world", prompt_mode=YoloePromptMode.LRPC), vis_module("foxglove"), - Agent.blueprint(), + McpServer.blueprint(), + McpClient.blueprint(), ).global_config(viewer="foxglove") From e382113444728e8b35edd99837deecf9fdddb717 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Apr 2026 20:15:01 -0700 Subject: [PATCH 33/53] =?UTF-8?q?fix(ci):=20regenerate=20all=5Fblueprints.?= =?UTF-8?q?py,=20fix=20test=5Fwebsocket=5Fserver=20stop=5Fexplore=5Fcmd=20?= =?UTF-8?q?=E2=86=92=20stop=5Fmovement?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dimos/robot/all_blueprints.py | 1 + dimos/visualization/rerun/test_websocket_server.py | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 00690d514f..b67f86b856 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -157,6 +157,7 @@ "reid-module": "dimos.perception.detection.reid.module.ReidModule", "replanning-a-star-planner": "dimos.navigation.replanning_a_star.module.ReplanningAStarPlanner", "rerun-bridge-module": "dimos.visualization.rerun.bridge.RerunBridgeModule", + "rerun-web-socket-server": "dimos.visualization.rerun.websocket_server.RerunWebSocketServer", "ros-nav": "dimos.navigation.rosnav.ROSNav", "simple-phone-teleop": "dimos.teleop.phone.phone_extensions.SimplePhoneTeleop", "spatial-memory": "dimos.perception.spatial_perception.SpatialMemory", diff --git a/dimos/visualization/rerun/test_websocket_server.py b/dimos/visualization/rerun/test_websocket_server.py index cec85fbb11..7282caf458 100644 --- a/dimos/visualization/rerun/test_websocket_server.py +++ b/dimos/visualization/rerun/test_websocket_server.py @@ -371,8 +371,8 @@ def test_stop_publishes_zero_twist_on_tele_cmd_vel(self) -> None: tw = received[0] assert tw.is_zero() - def test_twist_publishes_stop_explore_cmd_on_first_twist(self) -> None: - """First twist publishes Bool(data=True) on stop_explore_cmd; stop resets.""" + def test_twist_publishes_stop_movement_on_first_twist(self) -> None: + """First twist publishes Bool(data=True) on stop_movement; stop resets.""" mod = _make_module() mod.start() _wait_for_server(_TEST_PORT) @@ -380,7 +380,7 @@ def test_twist_publishes_stop_explore_cmd_on_first_twist(self) -> None: explore_cmds: list[Any] = [] twists: list[Any] = [] first_done = threading.Event() - mod.stop_explore_cmd.subscribe(_collect(explore_cmds, first_done)) + mod.stop_movement.subscribe(_collect(explore_cmds, first_done)) with MockViewerPublisher(f"ws://127.0.0.1:{_TEST_PORT}/ws") as pub: pub.send_twist(0.5, 0.0, 0.0, 0.0, 0.0, 0.0) @@ -390,7 +390,7 @@ def test_twist_publishes_stop_explore_cmd_on_first_twist(self) -> None: assert len(explore_cmds) == 1 assert explore_cmds[0].data is True - # Second twist within same connection should NOT publish another stop_explore_cmd + # Second twist within same connection should NOT publish another stop_movement twist_done = threading.Event() mod.tele_cmd_vel.subscribe(_collect(twists, twist_done)) @@ -400,7 +400,7 @@ def test_twist_publishes_stop_explore_cmd_on_first_twist(self) -> None: twist_done.wait(timeout=2.0) assert len(explore_cmds) == 1 # still just the first one - # After stop + new twist within same connection, stop_explore_cmd should fire again + # After stop + new twist within same connection, stop_movement should fire again second_done = threading.Event() def _on_second(msg: Any) -> None: @@ -408,7 +408,7 @@ def _on_second(msg: Any) -> None: if len(explore_cmds) >= 2: second_done.set() - mod.stop_explore_cmd.subscribe(_on_second) + mod.stop_movement.subscribe(_on_second) pub.send_stop() pub.send_twist(0.1, 0.0, 0.0, 0.0, 0.0, 0.0) From 8e1d71aadd32c9c55509d3a7468dddf263b5e98e Mon Sep 17 00:00:00 2001 From: stash Date: Thu, 2 Apr 2026 11:53:42 -0700 Subject: [PATCH 34/53] Revered topics to cmd_vel to fix unitree go2 blueprints and fallback to connect to grpc if port already used --- dimos/visualization/rerun/bridge.py | 1 + dimos/visualization/rerun/websocket_server.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 843ae421f4..a562eaede3 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -363,6 +363,7 @@ def start(self) -> None: grpc_port=grpc_port, server_memory_limit=self.config.memory_limit, ) + rr.connect_grpc(self.config.connect_url) logger.info(f"Rerun gRPC server ready at {self.config.connect_url}") # "none" - just init, no viewer (connect externally) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index e7319d176a..5610f35c36 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -73,7 +73,7 @@ class RerunWebSocketServer(Module[Config]): default_config = Config clicked_point: Out[PointStamped] - tele_cmd_vel: Out[Twist] + cmd_vel: Out[Twist] stop_movement: Out[Bool] def __init__(self, **kwargs: Any) -> None: @@ -218,12 +218,12 @@ def _dispatch(self, raw: str | bytes, client_id: int) -> None: if not self._teleop_clients: self.stop_movement.publish(Bool(data=True)) self._teleop_clients.add(client_id) - self.tele_cmd_vel.publish(twist) + self.cmd_vel.publish(twist) elif msg_type == "stop": logger.debug("RerunWebSocketServer: stop") self._teleop_clients.discard(client_id) - self.tele_cmd_vel.publish(Twist.zero()) + self.cmd_vel.publish(Twist.zero()) elif msg_type == "heartbeat": logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") From 032756e35845e1edf58641427f4220095be4f6f8 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 13:46:51 -0700 Subject: [PATCH 35/53] new rerun cli options --- dimos/core/global_config.py | 7 +- dimos/robot/cli/dimos.py | 28 +++++- dimos/simulation/unity/blueprint.py | 4 +- dimos/visualization/constants.py | 9 ++ dimos/visualization/rerun/bridge.py | 139 ++++++++++------------------ dimos/visualization/vis_module.py | 14 ++- docs/usage/cli.md | 8 ++ docs/usage/visualization.md | 2 +- 8 files changed, 104 insertions(+), 107 deletions(-) create mode 100644 dimos/visualization/constants.py diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index ce08c74705..8ee6a7d147 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -13,15 +13,12 @@ # limitations under the License. import re -from typing import Literal, TypeAlias +from dimos.visualization.constants import ViewerBackend, RerunOpenOption, RERUN_OPEN_DEFAULT, RERUN_ENABLE_WEB from pydantic_settings import BaseSettings, SettingsConfigDict from dimos.models.vl.types import VlModelName -ViewerBackend: TypeAlias = Literal["rerun", "rerun-web", "rerun-connect", "foxglove", "none"] - - def _get_all_numbers(s: str) -> list[float]: return [float(x) for x in re.findall(r"-?\d+\.?\d*", s)] @@ -37,6 +34,8 @@ class GlobalConfig(BaseSettings): replay_dir: str = "go2_sf_office" new_memory: bool = False viewer: ViewerBackend = "rerun" + rerun_open: RerunOpenOption = RERUN_OPEN_DEFAULT + rerun_web: bool = RERUN_ENABLE_WEB n_workers: int = 2 memory_limit: str = "auto" mujoco_camera_position: str | None = None diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index 8a2be16668..af8b01c979 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -551,17 +551,35 @@ def send( @main.command(name="rerun-bridge") def rerun_bridge_cmd( - viewer_mode: str = typer.Option( - "native", help="Viewer mode: native (desktop), web (browser), none (headless)" - ), memory_limit: str = typer.Option( "25%", help="Memory limit for Rerun viewer (e.g., '4GB', '16GB', '25%')" ), + rerun_open: str = typer.Option( + "native", help="How to open Rerun: native, web, both, none" + ), + rerun_web: bool = typer.Option( + True, "--rerun-web/--no-rerun-web", help="Enable/Disable Rerun web server" + ), ) -> None: """Launch the Rerun visualization bridge.""" - from dimos.visualization.rerun.bridge import run_bridge + import signal + + from dimos.protocol.pubsub.impl.lcmpubsub import LCM + from dimos.protocol.service.lcmservice import autoconf + from dimos.visualization.rerun.bridge import RerunBridgeModule + + autoconf(check_only=True) + + bridge = RerunBridgeModule( + memory_limit=memory_limit, + rerun_open=rerun_open, + rerun_web=rerun_web, + pubsubs=[LCM()], + ) + bridge.start() - run_bridge(viewer_mode=viewer_mode, memory_limit=memory_limit) + signal.signal(signal.SIGINT, lambda *_: bridge.stop()) + signal.pause() if __name__ == "__main__": diff --git a/dimos/simulation/unity/blueprint.py b/dimos/simulation/unity/blueprint.py index 4dff253ca9..e6276ad1f5 100644 --- a/dimos/simulation/unity/blueprint.py +++ b/dimos/simulation/unity/blueprint.py @@ -28,7 +28,7 @@ from dimos.core.blueprints import autoconnect from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.simulation.unity.module import UnityBridgeModule -from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode +from dimos.visualization.rerun.bridge import RerunBridgeModule def _rerun_blueprint() -> Any: @@ -57,5 +57,5 @@ def _rerun_blueprint() -> Any: unity_sim = autoconnect( UnityBridgeModule.blueprint(), - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config), + RerunBridgeModule.blueprint(**rerun_config), ) diff --git a/dimos/visualization/constants.py b/dimos/visualization/constants.py new file mode 100644 index 0000000000..9d1239a085 --- /dev/null +++ b/dimos/visualization/constants.py @@ -0,0 +1,9 @@ +from typing import Literal, TypeAlias + +ViewerBackend: TypeAlias = Literal["rerun", "foxglove", "none"] +RerunOpenOption: TypeAlias = Literal["none", "web", "native", "both"] + +RERUN_OPEN_DEFAULT: RerunOpenOption = "native" +RERUN_ENABLE_WEB = True +RERUN_GRPC_PORT = 9876 +RERUN_WEB_PORT = 9090 \ No newline at end of file diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 843ae421f4..592849f1e9 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -23,20 +23,18 @@ import time from typing import ( Any, - Literal, Protocol, TypeAlias, TypeGuard, cast, runtime_checkable, + get_args, ) from reactivex.disposable import Disposable from rerun._baseclasses import Archetype from rerun.blueprint import Blueprint from toolz import pipe # type: ignore[import-untyped] -import typer - from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.msgs.sensor_msgs.Image import Image @@ -45,6 +43,7 @@ from dimos.protocol.pubsub.patterns import Glob, pattern_matches from dimos.protocol.pubsub.spec import SubscribeAllCapable from dimos.utils.logging_config import setup_logger +from dimos.visualization.constants import RerunOpenOption, RERUN_GRPC_PORT, RERUN_OPEN_DEFAULT, RERUN_ENABLE_WEB # Message types with large payloads that need rate-limiting. # Image (~1 MB/frame at 30 fps) and PointCloud2 (~600-800 KB/frame) @@ -53,9 +52,6 @@ # unthrottled so navigation overlays and user input are never dropped. _HEAVY_MSG_TYPES: tuple[type, ...] = (Image, PointCloud2) -RERUN_GRPC_PORT = 9876 -RERUN_WEB_PORT = 9090 - # TODO OUT visual annotations # @@ -128,9 +124,6 @@ class RerunConvertible(Protocol): def to_rerun(self) -> RerunData: ... -ViewerMode = Literal["native", "web", "connect", "none"] - - def _hex_to_rgba(hex_color: str) -> int: """Convert '#RRGGBB' to a 0xRRGGBBAA int (fully opaque).""" h = hex_color.lstrip("#") @@ -171,22 +164,6 @@ def _default_blueprint() -> Blueprint: ) -# Maps global_config.viewer -> bridge viewer_mode. -# Evaluated at blueprint construction time (main process), not in start() (worker process). -_BACKEND_TO_MODE: dict[str, ViewerMode] = { - "rerun": "native", - "rerun-web": "web", - "rerun-connect": "connect", - "none": "none", -} - - -def _resolve_viewer_mode() -> ViewerMode: - from dimos.core.global_config import global_config - - return _BACKEND_TO_MODE.get(global_config.viewer, "native") - - class Config(ModuleConfig): """Configuration for RerunBridgeModule.""" @@ -200,9 +177,10 @@ class Config(ModuleConfig): min_interval_sec: float = 0.1 # Rate-limit per entity path (default: 10 Hz max) entity_prefix: str = "world" topic_to_entity: Callable[[Any], str] | None = None - viewer_mode: ViewerMode = field(default_factory=_resolve_viewer_mode) connect_url: str = "rerun+http://127.0.0.1:9877/proxy" memory_limit: str = "25%" + rerun_open: RerunOpenOption = RERUN_OPEN_DEFAULT + rerun_web: bool = RERUN_ENABLE_WEB # Blueprint factory: callable(rrb) -> Blueprint for viewer layout configuration # Set to None to disable default blueprint @@ -312,17 +290,50 @@ def _on_message(self, msg: Any, topic: Any) -> None: @rpc def start(self) -> None: import rerun as rr + from urllib.parse import urlparse + import socket super().start() self._last_log: dict[str, float] = {} # reset on each start - logger.info("Rerun bridge starting", viewer_mode=self.config.viewer_mode) + logger.info("Rerun bridge starting") - # Initialize and spawn Rerun viewer + # Initialize rr.init("dimos") + + # start grpc if needed + # If the port is already in use (another instance running), connect - if self.config.viewer_mode == "native": - spawned = False + parsed = urlparse(self.config.connect_url.replace("rerun+", "", 1)) + grpc_port = parsed.port or RERUN_GRPC_PORT + + port_in_use = False + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + port_in_use = sock.connect_ex(("127.0.0.1", grpc_port)) == 0 + + if port_in_use: + logger.info( + f"gRPC port {grpc_port} already in use, connecting to existing server" + ) + rr.connect_grpc(url=self.config.connect_url) + server_uri = self.config.connect_url + else: + server_uri = rr.serve_grpc( + grpc_port=grpc_port, + server_memory_limit=self.config.memory_limit, + ) + logger.info(f"Rerun gRPC server ready at {server_uri}") + + # Check open arg + if self.config.rerun_open not in get_args(RerunOpenOption): + logger.warning( + f"rerun_open was {self.config.rerun_open} which is not one of {get_args(RerunOpenOption)}", + exc_info=True, + ) + + # launch native viewer if desired + spawned = False + if self.config.rerun_open == "native" or self.config.rerun_open == "both": try: import rerun_bindings @@ -339,6 +350,8 @@ def start(self) -> None: "dimos-viewer found but failed to spawn, falling back to stock rerun", exc_info=True, ) + + # fallback on normal (non-dimos-viewer) rerun if not spawned: try: rr.spawn(connect=True, memory_limit=self.config.memory_limit) @@ -349,23 +362,12 @@ def start(self) -> None: "accessible via rerun-connect or rerun-web.", exc_info=True, ) - elif self.config.viewer_mode == "web": - server_uri = rr.serve_grpc() - rr.serve_web_viewer(connect_to=server_uri, open_browser=False) - elif self.config.viewer_mode == "connect": - # Serve gRPC so external viewers (dimos-viewer) can connect to us. - # Extract the port from the connect_url to match what viewers expect. - from urllib.parse import urlparse - - parsed = urlparse(self.config.connect_url.replace("rerun+", "", 1)) - grpc_port = parsed.port or RERUN_GRPC_PORT - rr.serve_grpc( - grpc_port=grpc_port, - server_memory_limit=self.config.memory_limit, - ) - logger.info(f"Rerun gRPC server ready at {self.config.connect_url}") - # "none" - just init, no viewer (connect externally) - + # web + open_web = self.config.rerun_open == "web" or self.config.rerun_open == "both" + if open_web or self.config.rerun_web: + rr.serve_web_viewer(connect_to=server_uri, open_browser=open_web) + + # setup blueprint if self.config.blueprint: rr.send_blueprint(_with_graph_tab(self.config.blueprint())) @@ -467,46 +469,3 @@ def stop(self) -> None: super().stop() -def run_bridge( - viewer_mode: str = "native", - memory_limit: str = "25%", -) -> None: - """Start a RerunBridgeModule with default LCM config and block until interrupted.""" - import signal - - from dimos.protocol.service.lcmservice import autoconf - - autoconf(check_only=True) - - bridge = RerunBridgeModule( - viewer_mode=viewer_mode, - memory_limit=memory_limit, - # any pubsub that supports subscribe_all and topic that supports str(topic) - # is acceptable here - pubsubs=[LCM()], - ) - - bridge.start() - - signal.signal(signal.SIGINT, lambda *_: bridge.stop()) - signal.pause() - - -app = typer.Typer() - - -@app.command() -def cli( - viewer_mode: str = typer.Option( - "native", help="Viewer mode: native (desktop), web (browser), none (headless)" - ), - memory_limit: str = typer.Option( - "25%", help="Memory limit for Rerun viewer (e.g., '4GB', '16GB', '25%')" - ), -) -> None: - """Rerun bridge for LCM messages.""" - run_bridge(viewer_mode=viewer_mode, memory_limit=memory_limit) - - -if __name__ == "__main__": - app() diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index 400a912ce4..2a3f725493 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -64,15 +64,19 @@ def vis_module( RerunWebSocketServer.blueprint(), WebsocketVisModule.blueprint(), ) - case "rerun" | "rerun-web" | "rerun-connect": + case "rerun": from dimos.protocol.pubsub.impl.lcmpubsub import LCM - from dimos.visualization.rerun.bridge import _BACKEND_TO_MODE, RerunBridgeModule + from dimos.visualization.rerun.bridge import RerunBridgeModule + from dimos.core.global_config import global_config - rerun_config = {**rerun_config} + rerun_config = {**rerun_config} # copy (avoid mutation) rerun_config.setdefault("pubsubs", [LCM()]) - viewer_mode = _BACKEND_TO_MODE.get(viewer_backend, "native") + rerun_config.setdefault("rerun_open", global_config.rerun_open) + rerun_config.setdefault("rerun_web", global_config.rerun_web) return autoconnect( - RerunBridgeModule.blueprint(viewer_mode=viewer_mode, **rerun_config), + RerunBridgeModule.blueprint( + **rerun_config, + ), RerunWebSocketServer.blueprint(), WebsocketVisModule.blueprint(), ) diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 38366f5daf..20855f1db8 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -320,6 +320,14 @@ rerun-bridge Also available as `dimos rerun-bridge`. +**Options:** + +| Option | Default | Description | +|--------|---------|-------------| +| `--memory-limit` | `25%` | Memory limit for Rerun viewer (e.g., `4GB`, `16GB`, `25%`) | +| `--rerun-open` | `native` | How to open Rerun: `native`, `web`, `both`, `none` | +| `--rerun-web` / `--no-rerun-web` | `--rerun-web` | Enable/Disable Rerun web server | + --- ## File Locations diff --git a/docs/usage/visualization.md b/docs/usage/visualization.md index 57ad460354..820e674867 100644 --- a/docs/usage/visualization.md +++ b/docs/usage/visualization.md @@ -73,7 +73,7 @@ from dimos.protocol.pubsub.impl.lcmpubsub import LCM camera_demo = autoconnect( CameraModule.blueprint(), RerunBridgeModule.blueprint( - viewer_mode="native", # native (desktop), web (browser), none (headless) + rerun_open="native", # native, web, both, none ), ) From bf2e94cbae5d87cd1a788d84713ce22a9dba3360 Mon Sep 17 00:00:00 2001 From: jeff-hykin <17692058+jeff-hykin@users.noreply.github.com> Date: Thu, 2 Apr 2026 20:54:06 +0000 Subject: [PATCH 36/53] CI code cleanup --- dimos/core/global_config.py | 8 +++++++- dimos/robot/cli/dimos.py | 4 +--- dimos/visualization/constants.py | 16 ++++++++++++++- dimos/visualization/rerun/bridge.py | 31 ++++++++++++++++------------- dimos/visualization/vis_module.py | 4 ++-- 5 files changed, 42 insertions(+), 21 deletions(-) diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index 8ee6a7d147..7584efed04 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -13,11 +13,17 @@ # limitations under the License. import re -from dimos.visualization.constants import ViewerBackend, RerunOpenOption, RERUN_OPEN_DEFAULT, RERUN_ENABLE_WEB from pydantic_settings import BaseSettings, SettingsConfigDict from dimos.models.vl.types import VlModelName +from dimos.visualization.constants import ( + RERUN_ENABLE_WEB, + RERUN_OPEN_DEFAULT, + RerunOpenOption, + ViewerBackend, +) + def _get_all_numbers(s: str) -> list[float]: return [float(x) for x in re.findall(r"-?\d+\.?\d*", s)] diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index af8b01c979..deabedc10e 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -554,9 +554,7 @@ def rerun_bridge_cmd( memory_limit: str = typer.Option( "25%", help="Memory limit for Rerun viewer (e.g., '4GB', '16GB', '25%')" ), - rerun_open: str = typer.Option( - "native", help="How to open Rerun: native, web, both, none" - ), + rerun_open: str = typer.Option("native", help="How to open Rerun: native, web, both, none"), rerun_web: bool = typer.Option( True, "--rerun-web/--no-rerun-web", help="Enable/Disable Rerun web server" ), diff --git a/dimos/visualization/constants.py b/dimos/visualization/constants.py index 9d1239a085..3d22457033 100644 --- a/dimos/visualization/constants.py +++ b/dimos/visualization/constants.py @@ -1,3 +1,17 @@ +# 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. + from typing import Literal, TypeAlias ViewerBackend: TypeAlias = Literal["rerun", "foxglove", "none"] @@ -6,4 +20,4 @@ RERUN_OPEN_DEFAULT: RerunOpenOption = "native" RERUN_ENABLE_WEB = True RERUN_GRPC_PORT = 9876 -RERUN_WEB_PORT = 9090 \ No newline at end of file +RERUN_WEB_PORT = 9090 diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 592849f1e9..4c19ab8549 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -27,14 +27,15 @@ TypeAlias, TypeGuard, cast, - runtime_checkable, get_args, + runtime_checkable, ) from reactivex.disposable import Disposable from rerun._baseclasses import Archetype from rerun.blueprint import Blueprint from toolz import pipe # type: ignore[import-untyped] + from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.msgs.sensor_msgs.Image import Image @@ -43,7 +44,12 @@ from dimos.protocol.pubsub.patterns import Glob, pattern_matches from dimos.protocol.pubsub.spec import SubscribeAllCapable from dimos.utils.logging_config import setup_logger -from dimos.visualization.constants import RerunOpenOption, RERUN_GRPC_PORT, RERUN_OPEN_DEFAULT, RERUN_ENABLE_WEB +from dimos.visualization.constants import ( + RERUN_ENABLE_WEB, + RERUN_GRPC_PORT, + RERUN_OPEN_DEFAULT, + RerunOpenOption, +) # Message types with large payloads that need rate-limiting. # Image (~1 MB/frame at 30 fps) and PointCloud2 (~600-800 KB/frame) @@ -289,9 +295,10 @@ def _on_message(self, msg: Any, topic: Any) -> None: @rpc def start(self) -> None: - import rerun as rr - from urllib.parse import urlparse import socket + from urllib.parse import urlparse + + import rerun as rr super().start() @@ -300,7 +307,7 @@ def start(self) -> None: # Initialize rr.init("dimos") - + # start grpc if needed # If the port is already in use (another instance running), connect @@ -312,9 +319,7 @@ def start(self) -> None: port_in_use = sock.connect_ex(("127.0.0.1", grpc_port)) == 0 if port_in_use: - logger.info( - f"gRPC port {grpc_port} already in use, connecting to existing server" - ) + logger.info(f"gRPC port {grpc_port} already in use, connecting to existing server") rr.connect_grpc(url=self.config.connect_url) server_uri = self.config.connect_url else: @@ -323,14 +328,14 @@ def start(self) -> None: server_memory_limit=self.config.memory_limit, ) logger.info(f"Rerun gRPC server ready at {server_uri}") - + # Check open arg if self.config.rerun_open not in get_args(RerunOpenOption): logger.warning( f"rerun_open was {self.config.rerun_open} which is not one of {get_args(RerunOpenOption)}", exc_info=True, ) - + # launch native viewer if desired spawned = False if self.config.rerun_open == "native" or self.config.rerun_open == "both": @@ -350,7 +355,7 @@ def start(self) -> None: "dimos-viewer found but failed to spawn, falling back to stock rerun", exc_info=True, ) - + # fallback on normal (non-dimos-viewer) rerun if not spawned: try: @@ -366,7 +371,7 @@ def start(self) -> None: open_web = self.config.rerun_open == "web" or self.config.rerun_open == "both" if open_web or self.config.rerun_web: rr.serve_web_viewer(connect_to=server_uri, open_browser=open_web) - + # setup blueprint if self.config.blueprint: rr.send_blueprint(_with_graph_tab(self.config.blueprint())) @@ -467,5 +472,3 @@ def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: def stop(self) -> None: self._visual_override_for_entity_path.cache_clear() super().stop() - - diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index 2a3f725493..e9cac119c6 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -65,11 +65,11 @@ def vis_module( WebsocketVisModule.blueprint(), ) case "rerun": + from dimos.core.global_config import global_config from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.visualization.rerun.bridge import RerunBridgeModule - from dimos.core.global_config import global_config - rerun_config = {**rerun_config} # copy (avoid mutation) + rerun_config = {**rerun_config} # copy (avoid mutation) rerun_config.setdefault("pubsubs", [LCM()]) rerun_config.setdefault("rerun_open", global_config.rerun_open) rerun_config.setdefault("rerun_web", global_config.rerun_web) From 56a527faa12713df9f6420f1fe95c562ae340622 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 16:35:20 -0700 Subject: [PATCH 37/53] add connect printout for headless --- dimos/visualization/rerun/bridge.py | 37 ++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 4c19ab8549..9545cd226d 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -342,10 +342,12 @@ def start(self) -> None: try: import rerun_bindings + # Use --connect so the viewer connects to the bridge's gRPC + # server rather than starting its own (which would conflict). rerun_bindings.spawn( - port=RERUN_GRPC_PORT, executable_name="dimos-viewer", memory_limit=self.config.memory_limit, + extra_args=["--connect", server_uri], ) spawned = True except ImportError: @@ -360,6 +362,7 @@ def start(self) -> None: if not spawned: try: rr.spawn(connect=True, memory_limit=self.config.memory_limit) + spawned = True except (RuntimeError, FileNotFoundError): logger.warning( "Rerun native viewer not available (headless?). " @@ -367,10 +370,15 @@ def start(self) -> None: "accessible via rerun-connect or rerun-web.", exc_info=True, ) + # web open_web = self.config.rerun_open == "web" or self.config.rerun_open == "both" if open_web or self.config.rerun_web: rr.serve_web_viewer(connect_to=server_uri, open_browser=open_web) + + # printout + if self.config.rerun_open == "none" or (self.config.rerun_open == "native" and not spawned): + self._log_connect_hints(grpc_port) # setup blueprint if self.config.blueprint: @@ -391,6 +399,33 @@ def start(self) -> None: self._log_static() + def _log_connect_hints(self, grpc_port: int) -> None: + """Log CLI commands for connecting a viewer to this bridge.""" + import socket + + from dimos.utils.generic import get_local_ips + + local_ips = get_local_ips() + hostname = socket.gethostname() + connect_url = f"rerun+http://127.0.0.1:{grpc_port}/proxy" + + lines = [ + "", + "=" * 60, + "Rerun gRPC server running (no viewer opened)", + "", + "Connect a viewer:", + f" dimos-viewer --connect {connect_url}", + ] + for ip, iface in local_ips: + lines.append(f" dimos-viewer --connect rerun+http://{ip}:{grpc_port}/proxy # {iface}") + lines.append("") + lines.append(f" hostname: {hostname}") + lines.append("=" * 60) + lines.append("") + + logger.info("\n".join(lines)) + def _log_static(self) -> None: import rerun as rr From af325dc5e1a0da60dbf39a6912aa22f4c94c0202 Mon Sep 17 00:00:00 2001 From: jeff-hykin <17692058+jeff-hykin@users.noreply.github.com> Date: Thu, 2 Apr 2026 23:36:27 +0000 Subject: [PATCH 38/53] CI code cleanup --- dimos/visualization/rerun/bridge.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 9545cd226d..b1a86f9826 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -370,12 +370,12 @@ def start(self) -> None: "accessible via rerun-connect or rerun-web.", exc_info=True, ) - + # web open_web = self.config.rerun_open == "web" or self.config.rerun_open == "both" if open_web or self.config.rerun_web: rr.serve_web_viewer(connect_to=server_uri, open_browser=open_web) - + # printout if self.config.rerun_open == "none" or (self.config.rerun_open == "native" and not spawned): self._log_connect_hints(grpc_port) From 231d47b0bd529f060e7af9b63ec2a3c9a57802a0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 18:35:40 -0700 Subject: [PATCH 39/53] restore name to tele_cmd_vel --- dimos/visualization/rerun/websocket_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 5610f35c36..09cac8fd2c 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -73,7 +73,7 @@ class RerunWebSocketServer(Module[Config]): default_config = Config clicked_point: Out[PointStamped] - cmd_vel: Out[Twist] + tele_cmd_vel: Out[Twist] stop_movement: Out[Bool] def __init__(self, **kwargs: Any) -> None: From badc59d20ca21afb3158fac5b0c4290324e294cb Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 18:35:51 -0700 Subject: [PATCH 40/53] clean _resolve_viewer_mode --- .../unitree/go2/blueprints/agentic/unitree_go2_security.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_security.py b/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_security.py index f8b292b103..53e877c38e 100644 --- a/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_security.py +++ b/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_security.py @@ -18,7 +18,7 @@ from dimos.core.blueprints import autoconnect from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic import unitree_go2_agentic -from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode +from dimos.visualization.rerun.bridge import RerunBridgeModule def _convert_camera_info(camera_info: Any) -> Any: @@ -90,7 +90,7 @@ def _go2_rerun_blueprint() -> Any: unitree_go2_security = autoconnect( unitree_go2_agentic, - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode(), **rerun_config), + RerunBridgeModule.blueprint(**rerun_config), ) __all__ = ["unitree_go2_security"] From 53c50a5b7cd6357bd87147e30a7d681e85f3e58f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:16:58 -0700 Subject: [PATCH 41/53] wire in proper cmd_vel mux and cancel for go2 --- dimos/navigation/cmd_vel_mux.py | 124 ++++++++++++++++++ .../wavefront_frontier_goal_selector.py | 11 ++ dimos/navigation/replanning_a_star/module.py | 18 ++- .../unitree_go2_webrtc_keyboard_teleop.py | 4 + .../go2/blueprints/smart/unitree_go2.py | 2 + dimos/robot/unitree/keyboard_teleop.py | 10 +- dimos/visualization/rerun/websocket_server.py | 4 +- .../web/websocket_vis/websocket_vis_module.py | 6 +- 8 files changed, 167 insertions(+), 12 deletions(-) create mode 100644 dimos/navigation/cmd_vel_mux.py diff --git a/dimos/navigation/cmd_vel_mux.py b/dimos/navigation/cmd_vel_mux.py new file mode 100644 index 0000000000..f9c5752a7e --- /dev/null +++ b/dimos/navigation/cmd_vel_mux.py @@ -0,0 +1,124 @@ +# 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. + +"""CmdVelMux: merges nav and teleop velocity commands. + +Teleop (tele_cmd_vel) takes priority over autonomous navigation +(nav_cmd_vel). When teleop is active, nav commands are suppressed +and the planner's goal is cancelled. After a cooldown period with +no teleop input, nav commands resume. +""" + +from __future__ import annotations + +import threading +from typing import Any + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.navigation.replanning_a_star.module_spec import ReplanningAStarPlannerSpec +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class CmdVelMuxConfig(ModuleConfig): + teleop_cooldown_sec: float = 1.0 + + +class CmdVelMux(Module[CmdVelMuxConfig]): + """Multiplexes nav_cmd_vel and tele_cmd_vel into a single cmd_vel output. + + When teleop input arrives, the planner's current goal is cancelled + so the robot responds immediately to manual control. + + Ports: + nav_cmd_vel (In[Twist]): Velocity from the autonomous planner. + tele_cmd_vel (In[Twist]): Velocity from keyboard/joystick teleop. + cmd_vel (Out[Twist]): Merged output — teleop wins when active. + """ + + default_config = CmdVelMuxConfig + + nav_cmd_vel: In[Twist] + tele_cmd_vel: In[Twist] + cmd_vel: Out[Twist] + + _planner: ReplanningAStarPlannerSpec + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._teleop_active = False + self._lock = threading.Lock() + self._timer: threading.Timer | None = None + + def __getstate__(self) -> dict[str, Any]: + state = super().__getstate__() + state.pop("_lock", None) + state.pop("_timer", None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + super().__setstate__(state) + self._lock = threading.Lock() + self._timer = None + + @rpc + def start(self) -> None: + self.nav_cmd_vel._transport.subscribe(self._on_nav) + self.tele_cmd_vel._transport.subscribe(self._on_teleop) + + @rpc + def stop(self) -> None: + with self._lock: + if self._timer is not None: + self._timer.cancel() + self._timer = None + super().stop() + + def _on_nav(self, msg: Twist) -> None: + with self._lock: + if self._teleop_active: + return + self.cmd_vel._transport.publish(msg) + + def _on_teleop(self, msg: Twist) -> None: + was_active: bool + with self._lock: + was_active = self._teleop_active + self._teleop_active = True + if self._timer is not None: + self._timer.cancel() + self._timer = threading.Timer( + self.config.teleop_cooldown_sec, + self._end_teleop, + ) + self._timer.daemon = True + self._timer.start() + + if not was_active: + try: + self._planner.cancel_goal() + logger.info("Teleop active — cancelled planner goal") + except Exception: + logger.debug("Could not cancel planner goal", exc_info=True) + + self.cmd_vel._transport.publish(msg) + + def _end_teleop(self) -> None: + with self._lock: + self._teleop_active = False + self._timer = None diff --git a/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py b/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py index 4f5ade1a6f..56cc86dd5b 100644 --- a/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py +++ b/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py @@ -115,6 +115,7 @@ class WavefrontFrontierExplorer(Module[WavefrontConfig]): goal_reached: In[Bool] explore_cmd: In[Bool] stop_explore_cmd: In[Bool] + stop_movement: In[Bool] # LCM outputs goal_request: Out[PoseStamped] @@ -171,6 +172,10 @@ def start(self) -> None: unsub = self.stop_explore_cmd.subscribe(self._on_stop_explore_cmd) self._disposables.add(Disposable(unsub)) + if self.stop_movement.transport is not None: + unsub = self.stop_movement.subscribe(self._on_stop_movement) + self._disposables.add(Disposable(unsub)) + @rpc def stop(self) -> None: self.stop_exploration() @@ -201,6 +206,12 @@ def _on_stop_explore_cmd(self, msg: Bool) -> None: logger.info("Received exploration stop command via LCM") self.stop_exploration() + def _on_stop_movement(self, msg: Bool) -> None: + """Handle stop movement from teleop — cancel active exploration.""" + if msg.data and self.exploration_active: + logger.info("WavefrontFrontierExplorer: stop_movement received, stopping exploration") + self.stop_exploration() + def _count_costmap_information(self, costmap: OccupancyGrid) -> int: """ Count the amount of information in a costmap (free space + obstacles). diff --git a/dimos/navigation/replanning_a_star/module.py b/dimos/navigation/replanning_a_star/module.py index 26c540a254..5f9982a943 100644 --- a/dimos/navigation/replanning_a_star/module.py +++ b/dimos/navigation/replanning_a_star/module.py @@ -21,6 +21,9 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import In, Out +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() from dimos.msgs.geometry_msgs.PointStamped import PointStamped from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Twist import Twist @@ -36,10 +39,11 @@ class ReplanningAStarPlanner(Module, NavigationInterface): goal_request: In[PoseStamped] clicked_point: In[PointStamped] target: In[PoseStamped] + stop_movement: In[Bool] goal_reached: Out[Bool] navigation_state: Out[String] # TODO: set it - cmd_vel: Out[Twist] + nav_cmd_vel: Out[Twist] path: Out[Path] navigation_costmap: Out[OccupancyGrid] @@ -70,9 +74,14 @@ def start(self) -> None: ) ) + if self.stop_movement.transport is not None: + self._disposables.add( + Disposable(self.stop_movement.subscribe(self._on_stop_movement)) + ) + self._disposables.add(self._planner.path.subscribe(self.path.publish)) - self._disposables.add(self._planner.cmd_vel.subscribe(self.cmd_vel.publish)) + self._disposables.add(self._planner.cmd_vel.subscribe(self.nav_cmd_vel.publish)) self._disposables.add(self._planner.goal_reached.subscribe(self.goal_reached.publish)) @@ -90,6 +99,11 @@ def stop(self) -> None: super().stop() + def _on_stop_movement(self, msg: Bool) -> None: + if msg.data: + logger.info("ReplanningAStarPlanner: stop_movement received, cancelling goal") + self.cancel_goal() + @rpc def set_goal(self, goal: PoseStamped) -> bool: self._planner.handle_goal_request(goal) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_webrtc_keyboard_teleop.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_webrtc_keyboard_teleop.py index dad4054fa9..845a45e76f 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_webrtc_keyboard_teleop.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_webrtc_keyboard_teleop.py @@ -31,6 +31,10 @@ unitree_go2_webrtc_keyboard_teleop = autoconnect( unitree_go2_coordinator, KeyboardTeleop.blueprint(), +).remappings( + [ + (KeyboardTeleop, "tele_cmd_vel", "cmd_vel"), + ] ) __all__ = ["unitree_go2_webrtc_keyboard_teleop"] diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py index 194aff60ca..a90c9a4d9c 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py @@ -16,6 +16,7 @@ from dimos.core.blueprints import autoconnect from dimos.mapping.costmapper import CostMapper from dimos.mapping.voxels import VoxelGridMapper +from dimos.navigation.cmd_vel_mux import CmdVelMux from dimos.navigation.frontier_exploration.wavefront_frontier_goal_selector import ( WavefrontFrontierExplorer, ) @@ -30,6 +31,7 @@ ReplanningAStarPlanner.blueprint(), WavefrontFrontierExplorer.blueprint(), PatrollingModule.blueprint(), + CmdVelMux.blueprint(), ).global_config(n_workers=7, robot_model="unitree_go2") __all__ = ["unitree_go2"] diff --git a/dimos/robot/unitree/keyboard_teleop.py b/dimos/robot/unitree/keyboard_teleop.py index 4ebf6e3cce..23c2f312b2 100644 --- a/dimos/robot/unitree/keyboard_teleop.py +++ b/dimos/robot/unitree/keyboard_teleop.py @@ -33,10 +33,10 @@ class KeyboardTeleop(Module): """Pygame-based keyboard control module. - Outputs standard Twist messages on /cmd_vel for velocity control. + Outputs standard Twist messages on /tele_cmd_vel for velocity control. """ - cmd_vel: Out[Twist] # Standard velocity commands + tele_cmd_vel: Out[Twist] # Standard velocity commands _stop_event: threading.Event _keys_held: set[int] | None = None @@ -66,7 +66,7 @@ def stop(self) -> None: stop_twist = Twist() stop_twist.linear = Vector3(0, 0, 0) stop_twist.angular = Vector3(0, 0, 0) - self.cmd_vel.publish(stop_twist) + self.tele_cmd_vel.publish(stop_twist) self._stop_event.set() @@ -99,7 +99,7 @@ def _pygame_loop(self) -> None: stop_twist = Twist() stop_twist.linear = Vector3(0, 0, 0) stop_twist.angular = Vector3(0, 0, 0) - self.cmd_vel.publish(stop_twist) + self.tele_cmd_vel.publish(stop_twist) print("EMERGENCY STOP!") elif event.key == pygame.K_ESCAPE: # ESC quits @@ -143,7 +143,7 @@ def _pygame_loop(self) -> None: twist.angular.z *= speed_multiplier # Always publish twist at 50Hz - self.cmd_vel.publish(twist) + self.tele_cmd_vel.publish(twist) self._update_display(twist) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index 09cac8fd2c..e7319d176a 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -218,12 +218,12 @@ def _dispatch(self, raw: str | bytes, client_id: int) -> None: if not self._teleop_clients: self.stop_movement.publish(Bool(data=True)) self._teleop_clients.add(client_id) - self.cmd_vel.publish(twist) + self.tele_cmd_vel.publish(twist) elif msg_type == "stop": logger.debug("RerunWebSocketServer: stop") self._teleop_clients.discard(client_id) - self.cmd_vel.publish(Twist.zero()) + self.tele_cmd_vel.publish(Twist.zero()) elif msg_type == "heartbeat": logger.debug(f"RerunWebSocketServer: heartbeat ts={msg.get('timestamp_ms')}") diff --git a/dimos/web/websocket_vis/websocket_vis_module.py b/dimos/web/websocket_vis/websocket_vis_module.py index 39163bfb95..1075771576 100644 --- a/dimos/web/websocket_vis/websocket_vis_module.py +++ b/dimos/web/websocket_vis/websocket_vis_module.py @@ -104,7 +104,7 @@ class WebsocketVisModule(Module[WebsocketConfig]): gps_goal: Out[LatLon] explore_cmd: Out[Bool] stop_explore_cmd: Out[Bool] - cmd_vel: Out[Twist] + tele_cmd_vel: Out[Twist] movecmd_stamped: Out[TwistStamped] def __init__(self, **kwargs: Any) -> None: @@ -332,14 +332,14 @@ async def clear_gps_goals(sid: str) -> None: @self.sio.event # type: ignore[untyped-decorator] async def move_command(sid: str, data: dict[str, Any]) -> None: # Publish Twist if transport is configured - if self.cmd_vel and self.cmd_vel.transport: + if self.tele_cmd_vel and self.tele_cmd_vel.transport: twist = Twist( linear=Vector3(data["linear"]["x"], data["linear"]["y"], data["linear"]["z"]), angular=Vector3( data["angular"]["x"], data["angular"]["y"], data["angular"]["z"] ), ) - self.cmd_vel.publish(twist) + self.tele_cmd_vel.publish(twist) # Publish TwistStamped if transport is configured if self.movecmd_stamped and self.movecmd_stamped.transport: From f9ef7bcd356b4c2ea2f38c6562303049c86976ba Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:17:08 -0700 Subject: [PATCH 42/53] hide startup noise --- dimos/visualization/rerun/websocket_server.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/dimos/visualization/rerun/websocket_server.py b/dimos/visualization/rerun/websocket_server.py index e7319d176a..0be8a44bb2 100644 --- a/dimos/visualization/rerun/websocket_server.py +++ b/dimos/visualization/rerun/websocket_server.py @@ -31,6 +31,7 @@ import asyncio import json +import logging import threading from typing import Any @@ -150,6 +151,12 @@ async def _serve(self) -> None: self._stop_event = asyncio.Event() + # Suppress noisy tracebacks from non-WebSocket connections (e.g. port + # scanners, health checks, or accidental gRPC probes). The library + # logs failed handshakes at ERROR level, so we need CRITICAL to hide them. + ws_logger = logging.getLogger("websockets.server") + ws_logger.setLevel(logging.CRITICAL) + async with ws_server.serve( self._handle_client, host=self.config.host, @@ -158,6 +165,7 @@ async def _serve(self) -> None: # survive brief network hiccups while still detecting dead clients. ping_interval=30, ping_timeout=30, + logger=ws_logger, ): self._server_ready.set() await self._stop_event.wait() From a775501dda969dd6d452bac0d29233b362627d9c Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:17:33 -0700 Subject: [PATCH 43/53] get rid of warning --- dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 3efe8f4032..c164e3047b 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -65,7 +65,7 @@ def _static_base_link(rr: Any) -> list[Any]: rr.Boxes3D( half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)], - fill_mode="wireframe", + # fill_mode="wireframe", ), rr.Transform3D(parent_frame="tf#/base_link"), ] From 806acc94291bac60b5006a91f65e6c0ce451b43f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:34:29 -0700 Subject: [PATCH 44/53] fix sim on macos --- dimos/robot/unitree/mujoco_connection.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 03d15db756..7541819ea6 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -20,9 +20,12 @@ from collections.abc import Callable import functools import json +import os +from pathlib import Path import pickle import subprocess import sys +import sysconfig import threading import time from typing import Any, TypeVar @@ -126,12 +129,25 @@ def start(self) -> None: # Launch the subprocess try: - # mjpython must be used macOS (because of launch_passive inside mujoco_process.py) + # mjpython must be used on macOS (because of launch_passive inside mujoco_process.py). + # It needs libpython on the dylib search path; uv-installed Pythons + # use @rpath which doesn't always resolve inside venvs, so we + # point DYLD_LIBRARY_PATH at the real libpython directory. executable = sys.executable if sys.platform != "darwin" else "mjpython" + env = os.environ.copy() + if sys.platform == "darwin": + # on some systems mujoco looks in the wrong place for shared libraries. We + libdir = Path(sysconfig.get_config_var("LIBDIR") or "") + if libdir.is_dir(): + existing = env.get("DYLD_LIBRARY_PATH", "") + env["DYLD_LIBRARY_PATH"] = ( + f"{libdir}:{existing}" if existing else str(libdir) + ) self.process = subprocess.Popen( [executable, str(LAUNCHER_PATH), config_pickle, shm_names_json], stderr=subprocess.PIPE, + env=env, ) except Exception as e: From 71ce2a979d552395f4857ced915f7fe7bff29ecf Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:35:11 -0700 Subject: [PATCH 45/53] fix warning --- dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py | 1 - 1 file changed, 1 deletion(-) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index c164e3047b..8295b1d8ee 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -65,7 +65,6 @@ def _static_base_link(rr: Any) -> list[Any]: rr.Boxes3D( half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)], - # fill_mode="wireframe", ), rr.Transform3D(parent_frame="tf#/base_link"), ] From 0eed0b75da6ffaae830412775132aff5e5f7f765 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:35:38 -0700 Subject: [PATCH 46/53] comment --- dimos/robot/unitree/mujoco_connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 7541819ea6..9495a5a2b2 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -136,7 +136,7 @@ def start(self) -> None: executable = sys.executable if sys.platform != "darwin" else "mjpython" env = os.environ.copy() if sys.platform == "darwin": - # on some systems mujoco looks in the wrong place for shared libraries. We + # on some systems mujoco looks in the wrong place for shared libraries. So we force it look in the right place libdir = Path(sysconfig.get_config_var("LIBDIR") or "") if libdir.is_dir(): existing = env.get("DYLD_LIBRARY_PATH", "") From a90ac024932f70bcfe79248a25d3fcddfe8111b3 Mon Sep 17 00:00:00 2001 From: jeff-hykin <17692058+jeff-hykin@users.noreply.github.com> Date: Fri, 3 Apr 2026 02:38:13 +0000 Subject: [PATCH 47/53] CI code cleanup --- dimos/navigation/replanning_a_star/module.py | 4 +--- dimos/robot/unitree/mujoco_connection.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/dimos/navigation/replanning_a_star/module.py b/dimos/navigation/replanning_a_star/module.py index 5f9982a943..1e00ff441e 100644 --- a/dimos/navigation/replanning_a_star/module.py +++ b/dimos/navigation/replanning_a_star/module.py @@ -75,9 +75,7 @@ def start(self) -> None: ) if self.stop_movement.transport is not None: - self._disposables.add( - Disposable(self.stop_movement.subscribe(self._on_stop_movement)) - ) + self._disposables.add(Disposable(self.stop_movement.subscribe(self._on_stop_movement))) self._disposables.add(self._planner.path.subscribe(self.path.publish)) diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 9495a5a2b2..5104799faa 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -140,9 +140,7 @@ def start(self) -> None: libdir = Path(sysconfig.get_config_var("LIBDIR") or "") if libdir.is_dir(): existing = env.get("DYLD_LIBRARY_PATH", "") - env["DYLD_LIBRARY_PATH"] = ( - f"{libdir}:{existing}" if existing else str(libdir) - ) + env["DYLD_LIBRARY_PATH"] = f"{libdir}:{existing}" if existing else str(libdir) self.process = subprocess.Popen( [executable, str(LAUNCHER_PATH), config_pickle, shm_names_json], From 9d2abd8924245cb0cf1ff08fbe8313dd850d92e9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 19:41:07 -0700 Subject: [PATCH 48/53] add test --- dimos/navigation/test_cmd_vel_mux.py | 57 ++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 dimos/navigation/test_cmd_vel_mux.py diff --git a/dimos/navigation/test_cmd_vel_mux.py b/dimos/navigation/test_cmd_vel_mux.py new file mode 100644 index 0000000000..d7bc696973 --- /dev/null +++ b/dimos/navigation/test_cmd_vel_mux.py @@ -0,0 +1,57 @@ +# 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. + +"""Tests for CmdVelMux teleop/nav priority switching.""" + +from __future__ import annotations + +from dimos.navigation.cmd_vel_mux import CmdVelMux + + +class TestCmdVelMux: + def test_teleop_initially_inactive(self) -> None: + mux = CmdVelMux.__new__(CmdVelMux) + mux.__dict__["_teleop_active"] = False + assert not mux._teleop_active + + def test_end_teleop_clears_flag(self) -> None: + import threading + + mux = CmdVelMux.__new__(CmdVelMux) + mux.__dict__["_teleop_active"] = True + mux.__dict__["_timer"] = None + mux.__dict__["_lock"] = threading.Lock() + mux._end_teleop() + assert not mux._teleop_active + + def test_nav_suppressed_when_teleop_active(self) -> None: + """When _teleop_active is True, _on_nav returns early (no publish).""" + import threading + + mux = CmdVelMux.__new__(CmdVelMux) + mux.__dict__["_teleop_active"] = True + mux.__dict__["_lock"] = threading.Lock() + # _on_nav should return before reaching cmd_vel._transport.publish + # If it didn't return early, it would crash since cmd_vel has no transport + from dimos.msgs.geometry_msgs.Twist import Twist + from dimos.msgs.geometry_msgs.Vector3 import Vector3 + + mux._on_nav(Twist(linear=Vector3(1, 0, 0), angular=Vector3(0, 0, 0))) + assert mux._teleop_active # Still active, nav was suppressed + + def test_cooldown_default(self) -> None: + from dimos.navigation.cmd_vel_mux import CmdVelMuxConfig + + config = CmdVelMuxConfig() + assert config.teleop_cooldown_sec == 1.0 From 50dcfebcc2a832416b450470b41cd735b4810e62 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 20:11:14 -0700 Subject: [PATCH 49/53] change how stop_movement works with cmd_vel --- dimos/navigation/cmd_vel_mux.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/dimos/navigation/cmd_vel_mux.py b/dimos/navigation/cmd_vel_mux.py index f9c5752a7e..2e3c6d4416 100644 --- a/dimos/navigation/cmd_vel_mux.py +++ b/dimos/navigation/cmd_vel_mux.py @@ -16,8 +16,8 @@ Teleop (tele_cmd_vel) takes priority over autonomous navigation (nav_cmd_vel). When teleop is active, nav commands are suppressed -and the planner's goal is cancelled. After a cooldown period with -no teleop input, nav commands resume. +and a stop_movement signal is published. After a cooldown period +with no teleop input, nav commands resume. """ from __future__ import annotations @@ -25,11 +25,12 @@ import threading from typing import Any +from dimos_lcm.std_msgs import Bool + from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.navigation.replanning_a_star.module_spec import ReplanningAStarPlannerSpec from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -42,13 +43,14 @@ class CmdVelMuxConfig(ModuleConfig): class CmdVelMux(Module[CmdVelMuxConfig]): """Multiplexes nav_cmd_vel and tele_cmd_vel into a single cmd_vel output. - When teleop input arrives, the planner's current goal is cancelled - so the robot responds immediately to manual control. + When teleop input arrives, stop_movement is published so downstream + modules (planner, explorer) can cancel their active goals. Ports: nav_cmd_vel (In[Twist]): Velocity from the autonomous planner. tele_cmd_vel (In[Twist]): Velocity from keyboard/joystick teleop. cmd_vel (Out[Twist]): Merged output — teleop wins when active. + stop_movement (Out[Bool]): Published when teleop begins. """ default_config = CmdVelMuxConfig @@ -56,8 +58,7 @@ class CmdVelMux(Module[CmdVelMuxConfig]): nav_cmd_vel: In[Twist] tele_cmd_vel: In[Twist] cmd_vel: Out[Twist] - - _planner: ReplanningAStarPlannerSpec + stop_movement: Out[Bool] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) @@ -110,11 +111,8 @@ def _on_teleop(self, msg: Twist) -> None: self._timer.start() if not was_active: - try: - self._planner.cancel_goal() - logger.info("Teleop active — cancelled planner goal") - except Exception: - logger.debug("Could not cancel planner goal", exc_info=True) + self.stop_movement.publish(Bool(data=True)) + logger.info("Teleop active — published stop_movement") self.cmd_vel._transport.publish(msg) From cdd11adeda92e8bb86debb3b61dd428292f8a7c5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 20:47:22 -0700 Subject: [PATCH 50/53] fix(mypy): resolve 7 type errors across 4 files - docker_module.py: import RERUN_GRPC_PORT/RERUN_WEB_PORT from constants (not bridge) - vis_module.py: import ViewerBackend from constants (not global_config) - cmd_vel_mux.py: add type: ignore for untyped __getstate__ super() call - websocket_vis_module.py: suppress comparison-overlap for rerun-web checks --- dimos/core/docker_module.py | 2 +- dimos/navigation/cmd_vel_mux.py | 2 +- dimos/visualization/vis_module.py | 2 +- dimos/web/websocket_vis/websocket_vis_module.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/dimos/core/docker_module.py b/dimos/core/docker_module.py index ae6b685312..424f903ea8 100644 --- a/dimos/core/docker_module.py +++ b/dimos/core/docker_module.py @@ -30,7 +30,7 @@ from dimos.core.rpc_client import ModuleProxyProtocol, RpcCall from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.utils.logging_config import setup_logger -from dimos.visualization.rerun.bridge import RERUN_GRPC_PORT, RERUN_WEB_PORT +from dimos.visualization.constants import RERUN_GRPC_PORT, RERUN_WEB_PORT if TYPE_CHECKING: from collections.abc import Callable diff --git a/dimos/navigation/cmd_vel_mux.py b/dimos/navigation/cmd_vel_mux.py index 2e3c6d4416..cc3683fa56 100644 --- a/dimos/navigation/cmd_vel_mux.py +++ b/dimos/navigation/cmd_vel_mux.py @@ -67,7 +67,7 @@ def __init__(self, **kwargs: Any) -> None: self._timer: threading.Timer | None = None def __getstate__(self) -> dict[str, Any]: - state = super().__getstate__() + state: dict[str, Any] = super().__getstate__() # type: ignore[no-untyped-call] state.pop("_lock", None) state.pop("_timer", None) return state diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index e9cac119c6..1aec6a40c3 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -18,7 +18,7 @@ from typing import Any from dimos.core.blueprints import Blueprint, autoconnect -from dimos.core.global_config import ViewerBackend +from dimos.visualization.constants import ViewerBackend def vis_module( diff --git a/dimos/web/websocket_vis/websocket_vis_module.py b/dimos/web/websocket_vis/websocket_vis_module.py index 1075771576..b76205d49f 100644 --- a/dimos/web/websocket_vis/websocket_vis_module.py +++ b/dimos/web/websocket_vis/websocket_vis_module.py @@ -159,7 +159,7 @@ def start(self) -> None: # Auto-open browser only for rerun-web (dashboard with Rerun iframe + command center) # For rerun and foxglove, users access the command center manually if needed - if self.config.g.viewer == "rerun-web": + if self.config.g.viewer == "rerun-web": # type: ignore[comparison-overlap] url = f"http://localhost:{self.config.port}/" logger.info(f"Dimensional Command Center: {url}") @@ -236,7 +236,7 @@ def _create_server(self) -> None: async def serve_index(request): # type: ignore[no-untyped-def] """Serve appropriate HTML based on viewer mode.""" # If running native Rerun, redirect to standalone command center - if self.config.g.viewer != "rerun-web": + if self.config.g.viewer != "rerun-web": # type: ignore[comparison-overlap] return RedirectResponse(url="/command-center") # Otherwise serve full dashboard with Rerun iframe From 0ee0c89e275b90164739e41e7913dba27853ece3 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Apr 2026 20:48:51 -0700 Subject: [PATCH 51/53] fix(ci): regenerate all_blueprints.py (add cmd-vel-mux) --- dimos/robot/all_blueprints.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index baad42c9dd..dff4728a00 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -104,6 +104,7 @@ "b1-connection-module": "dimos.robot.unitree.b1.connection.B1ConnectionModule", "camera-module": "dimos.hardware.sensors.camera.module.CameraModule", "cartesian-motion-controller": "dimos.manipulation.control.servo_control.cartesian_motion_controller.CartesianMotionController", + "cmd-vel-mux": "dimos.navigation.cmd_vel_mux.CmdVelMux", "control-coordinator": "dimos.control.coordinator.ControlCoordinator", "cost-mapper": "dimos.mapping.costmapper.CostMapper", "demo-calculator-skill": "dimos.agents.skills.demo_calculator_skill.DemoCalculatorSkill", From 24ca9d850e496053d7ee7e88ed22ff81033320c1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 4 Apr 2026 17:35:42 -0700 Subject: [PATCH 52/53] cleanup --- dimos/navigation/cmd_vel_mux.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dimos/navigation/cmd_vel_mux.py b/dimos/navigation/cmd_vel_mux.py index cc3683fa56..ded2233b1a 100644 --- a/dimos/navigation/cmd_vel_mux.py +++ b/dimos/navigation/cmd_vel_mux.py @@ -79,8 +79,8 @@ def __setstate__(self, state: dict[str, Any]) -> None: @rpc def start(self) -> None: - self.nav_cmd_vel._transport.subscribe(self._on_nav) - self.tele_cmd_vel._transport.subscribe(self._on_teleop) + self.nav_cmd_vel.subscribe(self._on_nav) + self.tele_cmd_vel.subscribe(self._on_teleop) @rpc def stop(self) -> None: @@ -94,7 +94,7 @@ def _on_nav(self, msg: Twist) -> None: with self._lock: if self._teleop_active: return - self.cmd_vel._transport.publish(msg) + self.cmd_vel.publish(msg) def _on_teleop(self, msg: Twist) -> None: was_active: bool @@ -114,7 +114,7 @@ def _on_teleop(self, msg: Twist) -> None: self.stop_movement.publish(Bool(data=True)) logger.info("Teleop active — published stop_movement") - self.cmd_vel._transport.publish(msg) + self.cmd_vel.publish(msg) def _end_teleop(self) -> None: with self._lock: From 128660639850f9ced892c31f61b2df0c96dc75a1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Apr 2026 12:14:33 -0700 Subject: [PATCH 53/53] fix: update blueprints import path in vis_module --- dimos/visualization/vis_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index 1aec6a40c3..3f0d5f2fe2 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -17,7 +17,7 @@ from typing import Any -from dimos.core.blueprints import Blueprint, autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.visualization.constants import ViewerBackend