From 957c26e50763cfb21580f9024cec4c973f3f73bf Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 21 Mar 2026 15:19:41 -0700 Subject: [PATCH 01/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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()