From aa0544819a4c5955219c9da7992b2bdd059c940e Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 29 Mar 2026 16:49:24 +0100 Subject: [PATCH 01/33] Memory2 record/replay functionality --- .gitignore | 3 + dimos/core/blueprints.py | 98 ++++ dimos/core/global_config.py | 1 + dimos/core/module_coordinator.py | 8 +- dimos/memory2/backend.py | 12 + dimos/memory2/codecs/base.py | 4 +- dimos/memory2/observationstore/base.py | 4 + dimos/memory2/observationstore/memory.py | 8 + dimos/memory2/observationstore/sqlite.py | 16 + dimos/memory2/stream.py | 6 + dimos/protocol/pubsub/impl/lcmpubsub.py | 5 +- dimos/protocol/pubsub/spec.py | 2 +- dimos/record/__init__.py | 17 + dimos/record/record_replay.py | 389 ++++++++++++++++ dimos/record/test_record_replay.py | 255 +++++++++++ dimos/robot/cli/dimos.py | 57 ++- dimos/types/__init__.py | 0 dimos/utils/cli/recorder/__init__.py | 0 dimos/utils/cli/recorder/run_recorder.py | 542 +++++++++++++++++++++++ flake.nix | 1 + pyproject.toml | 2 + uv.lock | 2 + 22 files changed, 1402 insertions(+), 30 deletions(-) create mode 100644 dimos/record/__init__.py create mode 100644 dimos/record/record_replay.py create mode 100644 dimos/record/test_record_replay.py create mode 100644 dimos/types/__init__.py create mode 100644 dimos/utils/cli/recorder/__init__.py create mode 100644 dimos/utils/cli/recorder/run_recorder.py diff --git a/.gitignore b/.gitignore index 4045db012e..4c34c551f4 100644 --- a/.gitignore +++ b/.gitignore @@ -77,3 +77,6 @@ CLAUDE.MD htmlcov/ .coverage .coverage.* + +# Created from simulation +MUJOCO_LOG.TXT diff --git a/dimos/core/blueprints.py b/dimos/core/blueprints.py index cac8507881..fa18cbe83e 100644 --- a/dimos/core/blueprints.py +++ b/dimos/core/blueprints.py @@ -30,6 +30,8 @@ from dimos.core.module_coordinator import ModuleCoordinator from dimos.core.stream import In, Out from dimos.core.transport import LCMTransport, PubSubTransport, pLCMTransport +from dimos.protocol.pubsub.impl.lcmpubsub import LCM +from dimos.record.record_replay import RecordReplay from dimos.spec.utils import Spec, is_spec, spec_annotation_compliance, spec_structural_compliance from dimos.utils.generic import short_id from dimos.utils.logging_config import setup_logger @@ -471,6 +473,79 @@ def _connect_rpc_methods(self, module_coordinator: ModuleCoordinator) -> None: requested_method_name, rpc_methods_dot[requested_method_name] ) + def replay( + self, + recording: RecordReplay | str, + *, + speed: float = 1.0, + cli_config_overrides: Mapping[str, Any] | None = None, + ) -> ModuleCoordinator: + """Build the blueprint with a recording providing some module outputs. + + Modules whose OUT streams are fully covered by the recording are + disabled — their data comes from the recording instead. All other + modules run normally. + + Args: + recording: A :class:`RecordReplay` instance, or a str + to a ``.db`` recording file. + speed: Playback speed multiplier (1.0 = realtime). + cli_config_overrides: Extra global config overrides. + + Returns: + The running :class:`ModuleCoordinator`. + """ + if isinstance(recording, str): + recording = RecordReplay(recording) + + recorded_streams = set(recording.store.list_streams()) + if not recorded_streams: + raise ValueError("Recording is empty — no streams to replay") + + # Find modules whose OUTs overlap with the recording. + # If ANY OUTs are covered, disable the module — the recording + # replaces it. Uncovered OUTs (e.g. on SHM, or never published) + # are simply absent during replay; downstream modules that need + # them won't receive data, which is the expected degraded mode. + modules_to_disable: list[type[ModuleBase]] = [] + for bp in self.blueprints: + out_names = {conn.name for conn in bp.streams if conn.direction == "out"} + if not out_names: + continue + covered = out_names & recorded_streams + if covered: + modules_to_disable.append(bp.module) + uncovered = out_names - covered + if uncovered: + logger.warning( + "Replay: disabling %s (partial coverage: replaying %s, missing %s)", + bp.module.__name__, + covered, + uncovered, + ) + else: + logger.info( + "Replay: disabling %s (all OUTs covered)", + bp.module.__name__, + ) + + if not modules_to_disable: + logger.warning( + "Replay: no modules disabled — recording streams %s " + "don't match any module OUT names", + recorded_streams, + ) + + patched = self.disabled_modules(*modules_to_disable) + coordinator = patched.build(cli_config_overrides) + + # Start playback in background — publishes to LCM so other modules receive data + lcm = LCM() + lcm.start() + recording.play(pubsub=lcm, speed=speed) + + return coordinator + def build( self, cli_config_overrides: Mapping[str, Any] | None = None, @@ -480,6 +555,29 @@ def build( if cli_config_overrides: global_config.update(**dict(cli_config_overrides)) + # Auto-replay if --replay-file is set in global config + replay_file = global_config.replay_file + if replay_file: + logger.info("Auto-replay from %s", replay_file) + # Strip replay_file from all override sources so the nested + # build() inside replay() does not re-enter this branch. + global_config.replay_file = None + clean_cli = ( + {k: v for k, v in cli_config_overrides.items() if k != "replay_file"} + if cli_config_overrides + else None + ) + clean_bp = replace( + self, + global_config_overrides=MappingProxyType( + {k: v for k, v in self.global_config_overrides.items() if k != "replay_file"} + ), + ) + return clean_bp.replay( + replay_file, + cli_config_overrides=clean_cli, + ) + self._run_configurators() self._check_requirements() self._verify_no_name_conflicts() diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index 60072ae7fd..3e2c700d58 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -33,6 +33,7 @@ class GlobalConfig(BaseSettings): simulation: bool = False replay: bool = False replay_dir: str = "go2_sf_office" + replay_file: str | None = None new_memory: bool = False viewer: ViewerBackend = "rerun" n_workers: int = 2 diff --git a/dimos/core/module_coordinator.py b/dimos/core/module_coordinator.py index 10227eae93..e048c3cfd2 100644 --- a/dimos/core/module_coordinator.py +++ b/dimos/core/module_coordinator.py @@ -14,8 +14,8 @@ from __future__ import annotations +import asyncio from concurrent.futures import ThreadPoolExecutor -import threading from typing import TYPE_CHECKING, Any from dimos.core.global_config import GlobalConfig, global_config @@ -154,10 +154,10 @@ def start_all_modules(self) -> None: def get_instance(self, module: type[ModuleBase]) -> ModuleProxy: return self._deployed_modules.get(module) # type: ignore[return-value, no-any-return] - def loop(self) -> None: - stop = threading.Event() + async def loop(self) -> None: + stop = asyncio.Event() try: - stop.wait() + await stop.wait() except KeyboardInterrupt: return finally: diff --git a/dimos/memory2/backend.py b/dimos/memory2/backend.py index c861993de9..b990700d46 100644 --- a/dimos/memory2/backend.py +++ b/dimos/memory2/backend.py @@ -16,6 +16,7 @@ from __future__ import annotations +from contextlib import suppress from dataclasses import replace from typing import TYPE_CHECKING, Any, Generic, TypeVar @@ -220,6 +221,17 @@ def _iterate_live( finally: sub.dispose() + def delete_range(self, t1: float, t2: float) -> int: + """Delete observations in [t1, t2] from all stores. Returns count deleted.""" + ids = self.metadata_store.delete_range(t1, t2) + for obs_id in ids: + if self.blob_store is not None: + with suppress(KeyError): + self.blob_store.delete(self.name, obs_id) + if self.vector_store is not None: + self.vector_store.delete(self.name, obs_id) + return len(ids) + def count(self, query: StreamQuery) -> int: if query.search_vec: return sum(1 for _ in self.iterate(query)) diff --git a/dimos/memory2/codecs/base.py b/dimos/memory2/codecs/base.py index 821b36b60f..01447de852 100644 --- a/dimos/memory2/codecs/base.py +++ b/dimos/memory2/codecs/base.py @@ -17,6 +17,8 @@ import importlib from typing import Any, Protocol, TypeVar, runtime_checkable +from dimos.msgs.sensor_msgs.Image import Image + T = TypeVar("T") @@ -33,8 +35,6 @@ def codec_for(payload_type: type[Any] | None = None) -> Codec[Any]: from dimos.memory2.codecs.pickle import PickleCodec if payload_type is not None: - from dimos.msgs.sensor_msgs.Image import Image - if issubclass(payload_type, Image): from dimos.memory2.codecs.jpeg import JpegCodec diff --git a/dimos/memory2/observationstore/base.py b/dimos/memory2/observationstore/base.py index 4d94889fb0..abe2f50e8e 100644 --- a/dimos/memory2/observationstore/base.py +++ b/dimos/memory2/observationstore/base.py @@ -69,5 +69,9 @@ def fetch_by_ids(self, ids: list[int]) -> list[Observation[T]]: """Batch fetch by id (for vector search results).""" ... + def delete_range(self, t1: float, t2: float) -> list[int]: + """Delete observations with ts in [t1, t2]. Returns deleted IDs.""" + raise NotImplementedError + def serialize(self) -> dict[str, Any]: return {"class": qual(type(self)), "config": self.config.model_dump()} diff --git a/dimos/memory2/observationstore/memory.py b/dimos/memory2/observationstore/memory.py index 529cd06394..19c90eb109 100644 --- a/dimos/memory2/observationstore/memory.py +++ b/dimos/memory2/observationstore/memory.py @@ -78,3 +78,11 @@ def fetch_by_ids(self, ids: list[int]) -> list[Observation[T]]: id_set = set(ids) with self._lock: return [obs for obs in self._observations if obs.id in id_set] + + def delete_range(self, t1: float, t2: float) -> list[int]: + """Delete observations with ts in [t1, t2]. Returns deleted IDs.""" + with self._lock: + to_delete = [obs for obs in self._observations if t1 <= obs.ts <= t2] + ids = [obs.id for obs in to_delete] + self._observations = [obs for obs in self._observations if not (t1 <= obs.ts <= t2)] + return ids diff --git a/dimos/memory2/observationstore/sqlite.py b/dimos/memory2/observationstore/sqlite.py index 5d680c540a..0146117ea8 100644 --- a/dimos/memory2/observationstore/sqlite.py +++ b/dimos/memory2/observationstore/sqlite.py @@ -440,5 +440,21 @@ def fetch_by_ids(self, ids: list[int]) -> list[Observation[T]]: rows = self._conn.execute(sql, ids).fetchall() return [self._row_to_obs(r, has_blob=join) for r in rows] + def delete_range(self, t1: float, t2: float) -> list[int]: + """Delete observations with ts in [t1, t2]. Returns deleted IDs.""" + with self._lock: + rows = self._conn.execute( + f'SELECT id FROM "{self._name}" WHERE ts >= ? AND ts <= ?', (t1, t2) + ).fetchall() + ids = [r[0] for r in rows] + if ids: + placeholders = ",".join("?" * len(ids)) + self._conn.execute(f'DELETE FROM "{self._name}" WHERE id IN ({placeholders})', ids) + self._conn.execute( + f'DELETE FROM "{self._name}_rtree" WHERE id IN ({placeholders})', ids + ) + self._conn.commit() + return ids + def stop(self) -> None: super().stop() diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index 545d387c32..b3b828d3d7 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -335,6 +335,12 @@ def subscribe( on_completed=on_completed, ) + def delete_range(self, t1: float, t2: float) -> int: + """Delete all observations with timestamps in [t1, t2]. Returns count deleted.""" + if isinstance(self._source, Stream): + raise TypeError("Cannot delete from a transform stream.") + return self._source.delete_range(t1, t2) + def append( self, payload: T, diff --git a/dimos/protocol/pubsub/impl/lcmpubsub.py b/dimos/protocol/pubsub/impl/lcmpubsub.py index 50c7c49f2f..7933c2503f 100644 --- a/dimos/protocol/pubsub/impl/lcmpubsub.py +++ b/dimos/protocol/pubsub/impl/lcmpubsub.py @@ -18,7 +18,6 @@ from dataclasses import dataclass import re import threading -from typing import Any from dimos.msgs.protocol import DimosMsg from dimos.protocol.pubsub.encoders import ( @@ -73,7 +72,7 @@ def from_channel_str(channel: str, default_lcm_type: type[DimosMsg] | None = Non return Topic(topic=topic_str, lcm_type=lcm_type or default_lcm_type) -class LCMPubSubBase(LCMService, AllPubSub[Topic, Any]): +class LCMPubSubBase(LCMService, AllPubSub[Topic, bytes]): """LCM-based PubSub with native regex subscription support. LCM natively supports regex patterns in subscribe(), so we implement @@ -92,7 +91,7 @@ def publish(self, topic: Topic | str, message: bytes) -> None: topic_str = str(topic) if isinstance(topic, Topic) else topic self.l.publish(topic_str, message) - def subscribe_all(self, callback: Callable[[bytes, Topic], Any]) -> Callable[[], None]: + def subscribe_all(self, callback: Callable[[bytes, Topic], None]) -> Callable[[], None]: return self.subscribe(Topic(re.compile(".*")), callback) # type: ignore[arg-type] def subscribe( diff --git a/dimos/protocol/pubsub/spec.py b/dimos/protocol/pubsub/spec.py index fe979fce82..0e292cfcb7 100644 --- a/dimos/protocol/pubsub/spec.py +++ b/dimos/protocol/pubsub/spec.py @@ -186,6 +186,6 @@ class SubscribeAllCapable(Protocol[MsgT_co, TopicT_co]): Both AllPubSub (native) and DiscoveryPubSub (synthesized) satisfy this. """ - def subscribe_all(self, callback: Callable[[Any, Any], Any]) -> Callable[[], None]: + def subscribe_all(self, callback: Callable[[MsgT_co, TopicT_co], None]) -> Callable[[], None]: """Subscribe to all messages on all topics.""" ... diff --git a/dimos/record/__init__.py b/dimos/record/__init__.py new file mode 100644 index 0000000000..1bb0bfbcc9 --- /dev/null +++ b/dimos/record/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.record.record_replay import RecordReplay + +__all__ = ("RecordReplay",) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py new file mode 100644 index 0000000000..4ef31e7511 --- /dev/null +++ b/dimos/record/record_replay.py @@ -0,0 +1,389 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""RecordReplay — record and replay pub/sub topics using memory2 stores. + +Usage:: + + from dimos.record import RecordReplay + from dimos.protocol.pubsub.impl.lcmpubsub import LCM + + # Record from live LCM traffic + rec = RecordReplay(Path("my_recording.db")) + rec.start_recording([LCM()]) + # ... robot runs ... + rec.stop_recording() + + # Replay into LCM (viewable via rerun-bridge) + rec.play(pubsub=LCM(), speed=1.0) + + # Timeline editing + rec.trim(start=2.0, end=30.0) + rec.delete_range(start=10.0, end=12.0) + + # Query/filter via memory2 streams + rec.stream("lidar").time_range(0, 5).count() +""" + +import asyncio +from collections.abc import Callable, Collection, Container +from contextlib import suppress +import heapq +import logging +import math +import re +import sys +import time +from typing import Any, TypedDict + +from dimos.memory2.codecs.base import _resolve_payload_type +from dimos.memory2.store.sqlite import SqliteStore +from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic +from dimos.protocol.pubsub.spec import PubSub + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing import Any as Self + +logger = logging.getLogger(__name__) + +_SANITIZE_RE = re.compile(r"[^A-Za-z0-9_]") + + +def topic_to_stream_name(channel: str) -> str: + """Convert a raw LCM channel/topic pattern to a safe stream name.""" + name = channel.split("#")[0].lstrip("/") + name = _SANITIZE_RE.sub("_", name) + if name and name[0].isdigit(): + name = f"_{name}" + return name or "_unknown" + + +class StreamInfo(TypedDict): + name: str + count: int + start: float + end: float + duration: float + type: str + + +class RecordReplay: + """Record and replay pub/sub topics using memory2 stores. + + A recording is a single SQLite file containing one stream per topic. + Supports recording from any ``SubscribeAllCapable`` pubsub, playback + with speed control, seeking, and timeline editing (trim/delete). + + The underlying :class:`SqliteStore` is fully accessible for advanced + queries via :attr:`store`, :attr:`streams`, and :meth:`stream`. + """ + + def __init__(self, path: str) -> None: + self._store = SqliteStore(path=path) + + self._recording = False + self._unsubscribes: list[Callable[[], None]] = [] + self._topic_filter: Container[str] | None = None + + self._resume = asyncio.Event() + self._play_task: asyncio.Task | None = None + self._play_speed = 1.0 + self._position = 0.0 + self._pubsub = None + + @property + def store(self) -> SqliteStore: + """The underlying store.""" + return self._store + + @property + def path(self) -> str: + """Path to the recording file.""" + return self._store.config.path + + @property + def is_recording(self) -> bool: + return self._recording + + @property + def is_playing(self) -> bool: + return self._play_task is not None and not self._play_task.done() + + @property + def is_paused(self) -> bool: + return self.is_playing and not self._resume.is_set() + + def start_recording( + self, + pubsubs: Collection[LCMPubSubBase], + topic_filter: Container[str] | None = None, + ) -> None: + """Start recording messages from the given pubsubs. + + Each pubsub is subscribed via ``subscribe_all()``. Messages are + stored in per-topic streams with automatic codec selection. + + Args: + pubsubs: List of pubsubs to subscribe to. + topic_filter: If provided, only record topics whose sanitized + stream name is in this set. If ``None``, record everything. + """ + if self._recording: + raise RuntimeError("Already recording") + self._recording = True + self._topic_filter = topic_filter + + for pubsub in pubsubs: + pubsub.start() + unsub = pubsub.subscribe_all(self._on_message) + self._unsubscribes.append(unsub) + + logger.info("Recording started on %d pubsub(s)", len(pubsubs)) + + def stop_recording(self) -> None: + """Stop recording.""" + if not self._recording: + return + self._recording = False + for unsub in self._unsubscribes: + unsub() + self._unsubscribes.clear() + logger.info("Recording stopped") + + def _on_message(self, msg: bytes, topic: Topic) -> None: + """Handle incoming message during recording.""" + stream_name = topic_to_stream_name(topic.pattern) + + if self._topic_filter is not None and stream_name not in self._topic_filter: + return + + msg_type = type(msg) + + s = self._store.stream(stream_name, msg_type) + s.append(msg, ts=time.time()) + + @property + def duration(self) -> float: + """Total duration of the recording in seconds.""" + t_min, t_max = self.time_range + return t_max - t_min + + @property + def time_range(self) -> tuple[float, float]: + """Absolute (min_ts, max_ts) across all streams.""" + streams = self._store.list_streams() + if not streams: + return (0.0, 0.0) + t_min = math.inf + t_max = -math.inf + for name in streams: + s = self._store.stream(name) + if s.exists(): + t0, t1 = s.get_time_range() + t_min = min(t_min, t0) + t_max = max(t_max, t1) + if t_min is math.inf: + return (0.0, 0.0) + return (t_min, t_max) + + @property + def position(self) -> float: + """Current playback position in seconds from recording start.""" + return self._position + + def stream_info(self) -> tuple[StreamInfo, ...]: + """Return per-stream metadata: name, count, time range, type.""" + result = [] + for name in self._store.list_streams(): + s = self._store.stream(name) + info: StreamInfo = {"name": name, "count": s.count()} + if info["count"] > 0: + t0, t1 = s.get_time_range() + info["start"] = t0 + info["end"] = t1 + info["duration"] = t1 - t0 + # Get payload type from registry + reg = self._store._registry.get(name) + if reg: + info["type"] = reg.get("payload_module", "unknown") + result.append(info) + return tuple(result) + + def play( + self, + pubsub: PubSub[Any, Any] | None = None, + speed: float = 1.0, + ) -> None: + """Start playback in a background thread. + + Args: + pubsub: If provided, messages are published to this pubsub. + Use LCM() to make them visible to rerun-bridge. + speed: Playback speed multiplier (1.0 = realtime, 2.0 = 2x, etc). + """ + if self.is_playing: + raise RuntimeError("Already playing") + + self._play_speed = speed + self._pubsub = pubsub + # Set resume so playback starts, this is cleared to pause playback. + self._resume.set() + self._play_task = asyncio.create_task(self._playback_loop(pubsub)) + + async def _playback_loop(self, pubsub: PubSub[Any, Any] | None) -> None: + """Main playback loop — merges all streams by timestamp.""" + t_min, t_max = self.time_range + if t_min >= t_max: + return + + # Build iterators for each stream, starting from seek position + start_ts = t_min + self._position + heap: list[tuple[float, int, str, Any]] = [] + counter = 0 # tiebreaker for heapq + + for name in self._store.list_streams(): + s = self._store.stream(name) + it = iter(s.after(start_ts - 0.001)) + try: + obs = next(it) + except StopIteration: + continue + heapq.heappush(heap, (obs.ts, counter, name, (obs, it))) + counter += 1 + + if not heap: + return + + topic_map = {} if pubsub is None else self._build_topic_map(pubsub) + wall_start = time.monotonic() + rec_start = heap[0][0] # earliest observation timestamp + + while heap: + if not self._resume.is_set(): + pause_start = time.monotonic() + await self._resume.wait() + wall_start += time.monotonic() - pause_start + + ts, _, stream_name, (obs, it) = heapq.heappop(heap) + + # Wait for correct playback time + elapsed_rec = ts - rec_start + target_wall = wall_start + (elapsed_rec / self._play_speed) + sleep_time = target_wall - time.monotonic() + await asyncio.sleep(sleep_time) + + self._position = ts - t_min + if pubsub is not None and stream_name in topic_map: + pubsub.publish(topic_map[stream_name], obs.data) + + try: + next_obs = next(it) + except StopIteration: + continue + heapq.heappush(heap, (next_obs.ts, counter, stream_name, (next_obs, it))) + counter += 1 + + def _build_topic_map(self, pubsub: PubSub[Any, Any]) -> dict[str, Topic]: + """Build stream_name -> Topic mapping for publishing.""" + topic_map: dict[str, Topic] = {} + + for name in self._store.list_streams(): + reg = self._store._registry.get(name) + if reg is None: + continue + + payload_module = reg.get("payload_module") + lcm_type = None + if payload_module: + lcm_type = _resolve_payload_type(payload_module) + + topic_map[name] = Topic(f"/{name}", lcm_type=lcm_type) + + return topic_map + + def pause(self) -> None: + """Pause playback. Resume with :meth:`resume`.""" + self._resume.clear() + + def resume(self) -> None: + """Resume paused playback.""" + self._resume.set() + + async def stop_playback(self) -> None: + """Stop playback.""" + self._resume.set() + if self._play_task is not None: + self._play_task.cancel() + with suppress(asyncio.CancelledError): + await self._play_task + self._play_task = None + + def seek(self, seconds: float) -> None: + """Set playback position in seconds from recording start. + + Takes effect immediately if playing (restarts playback loop). + If not playing, sets the position for the next :meth:`play`. + """ + self._position = max(0.0, min(seconds, self.duration)) + if self.is_playing: + self.stop_playback() + assert self._pubsub is not None + self.play(pubsub=self._pubsub, speed=self._play_speed) + + def delete_range(self, start: float, end: float) -> int: + """Delete observations in [start, end] seconds from recording start. + + Returns total count of deleted observations across all streams. + """ + t_min, _ = self.time_range + abs_start = t_min + start + abs_end = t_min + end + + total = 0 + for name in self._store.list_streams(): + s = self._store.stream(name) + total += s.delete_range(abs_start, abs_end) + return total + + def trim(self, start: float, end: float) -> int: + """Keep only [start, end] seconds, delete everything else. + + Returns total count of deleted observations. + """ + t_min, t_max = self.time_range + total = 0 + if start > 0: + total += self.delete_range(0, start - 0.0001) + if end < (t_max - t_min): + total += self.delete_range(end + 0.0001, t_max - t_min + 1) + return total + + async def close(self) -> None: + """Stop recording/playback and close the store.""" + self.stop_recording() + await self.stop_playback() + self._store.stop() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *exc: Any) -> None: + await self.close() + + def __repr__(self) -> str: + streams = self._store.list_streams() + dur = self.duration + return f"RecordReplay({self._store.config.path!r}, streams={len(streams)}, duration={dur:.1f}s)" diff --git a/dimos/record/test_record_replay.py b/dimos/record/test_record_replay.py new file mode 100644 index 0000000000..a14d1ef3cd --- /dev/null +++ b/dimos/record/test_record_replay.py @@ -0,0 +1,255 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for RecordReplay.""" + +import asyncio +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +import threading +from typing import Any + +import pytest + +from dimos.record import RecordReplay + + +class FakeTopic: + """Minimal topic for testing.""" + + def __init__(self, name: str) -> None: + self.topic = name + self.lcm_type = None + + @property + def pattern(self) -> str: + return self.topic + + def __str__(self) -> str: + return self.topic + + +class FakePubSub: + """Minimal PubSub that supports subscribe_all for testing.""" + + def __init__(self) -> None: + self._subscribers: list[Callable[[Any, Any], Any]] = [] + self._lock = threading.Lock() + self._started = False + + def start(self) -> None: + self._started = True + + def stop(self) -> None: + self._started = False + + def publish(self, topic: Any, message: Any) -> None: + # Not needed for recording tests + pass + + def subscribe_all(self, callback: Callable[[Any, Any], Any]) -> Callable[[], None]: + with self._lock: + self._subscribers.append(callback) + + def unsub() -> None: + with self._lock: + with suppress(ValueError): + self._subscribers.remove(callback) + + return unsub + + def emit(self, topic_name: str, msg: Any) -> None: + """Test helper: simulate a message arriving.""" + topic = FakeTopic(topic_name) + with self._lock: + subs = list(self._subscribers) + for cb in subs: + cb(msg, topic) + + +class SimpleMsg: + """A simple test message (not LCM, uses pickle codec).""" + + def __init__(self, value: float) -> None: + self.value = value + + def __eq__(self, other: object) -> bool: + return isinstance(other, SimpleMsg) and self.value == other.value + + +@pytest.fixture +def tmp_db(tmp_path: Path) -> str: + return str(tmp_path / "test_recording.db") + + +class TestRecordReplay: + async def test_record_and_list_streams(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + assert rec.is_recording + + pubsub.emit("/sensor/lidar", SimpleMsg(1.0)) + pubsub.emit("/sensor/odom", SimpleMsg(2.0)) + pubsub.emit("/sensor/lidar", SimpleMsg(3.0)) + await asyncio.sleep(0.05) # let timestamps diverge + + rec.stop_recording() + assert not rec.is_recording + + streams = rec.store.list_streams() + assert "sensor_lidar" in streams + assert "sensor_odom" in streams + + async def test_record_and_query(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + + for i in range(10): + pubsub.emit("/data", SimpleMsg(float(i))) + await asyncio.sleep(0.01) + + rec.stop_recording() + + s = rec.store.stream("data") + assert s.count() == 10 + first = s.first() + assert isinstance(first.data, SimpleMsg) + assert first.data.value == 0.0 + + async def test_duration(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + pubsub.emit("/a", SimpleMsg(0.0)) + await asyncio.sleep(0.1) + pubsub.emit("/a", SimpleMsg(1.0)) + rec.stop_recording() + + assert rec.duration >= 0.05 # at least some duration + + async def test_stream_info(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + for i in range(5): + pubsub.emit("/sensor", SimpleMsg(float(i))) + await asyncio.sleep(0.01) + rec.stop_recording() + + infos = rec.stream_info() + assert len(infos) == 1 + assert infos[0]["name"] == "sensor" + assert infos[0]["count"] == 5 + + async def test_delete_range(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + for i in range(20): + pubsub.emit("/data", SimpleMsg(float(i))) + await asyncio.sleep(0.01) + rec.stop_recording() + + before = rec.store.stream("data").count() + assert before == 20 + + dur = rec.duration + # Delete middle third + deleted = rec.delete_range(dur / 3, 2 * dur / 3) + assert deleted > 0 + + after = rec.store.stream("data").count() + assert after < before + + async def test_trim(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + for i in range(30): + pubsub.emit("/data", SimpleMsg(float(i))) + await asyncio.sleep(0.01) + rec.stop_recording() + + before = rec.store.stream("data").count() + dur = rec.duration + # Trim to middle third + rec.trim(dur / 3, 2 * dur / 3) + + after = rec.store.stream("data").count() + assert after < before + + async def test_repr(self, tmp_db: str) -> None: + async with RecordReplay(tmp_db) as rec: + r = repr(rec) + assert "RecordReplay" in r + assert "streams=0" in r + + async def test_playback_runs(self, tmp_db: str) -> None: + """Test that playback task starts and finishes.""" + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + for i in range(5): + pubsub.emit("/data", SimpleMsg(float(i))) + await asyncio.sleep(0.01) + rec.stop_recording() + + rec.play(speed=100.0) # very fast + assert rec.is_playing + async with asyncio.timeout(0.1): + await rec._play_task + assert not rec.is_playing + + async def test_stop_playback(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + for i in range(100): + pubsub.emit("/data", SimpleMsg(float(i))) + await asyncio.sleep(0.005) + rec.stop_recording() + + rec.play(speed=0.1) # slow + await asyncio.sleep(0.1) + assert rec.is_playing + await rec.stop_playback() + assert not rec.is_playing + + async def test_seek(self, tmp_db: str) -> None: + pubsub = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([pubsub]) + for i in range(10): + pubsub.emit("/data", SimpleMsg(float(i))) + await asyncio.sleep(0.01) + rec.stop_recording() + + rec.seek(0.05) + assert rec.position == pytest.approx(0.05, abs=0.01) + + async def test_multiple_pubsubs(self, tmp_db: str) -> None: + ps1 = FakePubSub() + ps2 = FakePubSub() + async with RecordReplay(tmp_db) as rec: + rec.start_recording([ps1, ps2]) + ps1.emit("/from_ps1", SimpleMsg(1.0)) + ps2.emit("/from_ps2", SimpleMsg(2.0)) + rec.stop_recording() + + streams = rec.store.list_streams() + assert "from_ps1" in streams + assert "from_ps2" in streams diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index 1137a612f3..17614034ac 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio from datetime import datetime, timezone import inspect import json @@ -28,9 +29,21 @@ import typer from dimos.agents.mcp.mcp_adapter import McpAdapter, McpError +from dimos.core.blueprints import autoconnect from dimos.core.global_config import GlobalConfig, global_config -from dimos.core.run_registry import get_most_recent, is_pid_alive, stop_entry -from dimos.utils.logging_config import setup_logger +from dimos.core.run_registry import ( + LOG_BASE_DIR, + RunEntry, + check_port_conflicts, + cleanup_stale, + generate_run_id, + get_most_recent, + is_pid_alive, + stop_entry, +) +from dimos.robot.get_all_blueprints import get_by_name, get_module_by_name +from dimos.utils.cli.recorder.run_recorder import main as recorder_main +from dimos.utils.logging_config import set_run_log_dir, setup_exception_handler, setup_logger logger = setup_logger() @@ -108,27 +121,12 @@ def callback(**kwargs) -> None: # type: ignore[no-untyped-def] main.callback()(create_dynamic_callback()) # type: ignore[no-untyped-call] -@main.command() -def run( +async def _run( ctx: typer.Context, robot_types: list[str] = typer.Argument(..., help="Blueprints or modules to run"), daemon: bool = typer.Option(False, "--daemon", "-d", help="Run in background"), disable: list[str] = typer.Option([], "--disable", help="Module names to disable"), ) -> None: - """Start a robot blueprint""" - logger.info("Starting DimOS") - - from dimos.core.blueprints import autoconnect - from dimos.core.run_registry import ( - LOG_BASE_DIR, - RunEntry, - check_port_conflicts, - cleanup_stale, - generate_run_id, - ) - from dimos.robot.get_all_blueprints import get_by_name, get_module_by_name - from dimos.utils.logging_config import set_run_log_dir, setup_exception_handler - setup_exception_handler() cli_config_overrides: dict[str, Any] = ctx.obj @@ -203,7 +201,7 @@ def run( ) entry.save() install_signal_handlers(entry, coordinator) - coordinator.loop() + await coordinator.loop() else: entry = RunEntry( run_id=run_id, @@ -217,11 +215,23 @@ def run( ) entry.save() try: - coordinator.loop() + await coordinator.loop() finally: entry.remove() +@main.command() +def run( + ctx: typer.Context, + robot_types: list[str] = typer.Argument(..., help="Blueprints or modules to run"), + daemon: bool = typer.Option(False, "--daemon", "-d", help="Run in background"), + disable: list[str] = typer.Option([], "--disable", help="Module names to disable"), +) -> None: + """Start a robot blueprint""" + logger.info("Starting DimOS") + asyncio.run(_run(ctx, robot_types, daemon, disable)) + + @main.command() def status() -> None: """Show the running DimOS instance.""" @@ -523,6 +533,13 @@ def top(ctx: typer.Context) -> None: dtop_main() +@main.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) +def recorder(ctx: typer.Context) -> None: + """Record and replay tool — terminal VLC for dimos recordings.""" + sys.argv = ["recorder", *ctx.args] + recorder_main() + + topic_app = typer.Typer(help="Topic commands for pub/sub") main.add_typer(topic_app, name="topic") diff --git a/dimos/types/__init__.py b/dimos/types/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dimos/utils/cli/recorder/__init__.py b/dimos/utils/cli/recorder/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py new file mode 100644 index 0000000000..4827abf6f2 --- /dev/null +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -0,0 +1,542 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""recorder — Terminal VLC for dimos recordings. + +Record from live LCM traffic, play back recordings, trim, seek. +Run ``rerun-bridge`` in another terminal to visualize playback. + +Usage:: + + recorder # interactive — record from LCM + recorder play my_recording.db # play an existing recording + recorder --help +""" + +from __future__ import annotations + +from collections import deque +import time +from typing import Any + +from rich.text import Text +from textual.app import App, ComposeResult +from textual.color import Color +from textual.containers import Horizontal +from textual.widgets import DataTable, Footer, Header, Static + +from dimos.utils.cli import theme + +# Braille sparkline constants (same as dtop) +_BRAILLE_BASE = 0x2800 +_LDOTS = (0x00, 0x40, 0x44, 0x46, 0x47) +_RDOTS = (0x00, 0x80, 0xA0, 0xB0, 0xB8) +_SPARK_WIDTH = 16 + +from dimos.record.record_replay import topic_to_stream_name + + +def _heat(ratio: float) -> str: + """Map 0..1 to cyan -> yellow -> red.""" + cyan = Color.parse(theme.CYAN) + yellow = Color.parse(theme.YELLOW) + red = Color.parse(theme.RED) + if ratio <= 0.5: + return cyan.blend(yellow, ratio * 2).hex + return yellow.blend(red, (ratio - 0.5) * 2).hex + + +def _spark(history: deque[float], max_val: float, width: int = _SPARK_WIDTH) -> Text: + """Braille sparkline from a deque of values.""" + n = width * 2 + vals = list(history) + if len(vals) < n: + vals = [0.0] * (n - len(vals)) + vals + else: + vals = vals[-n:] + result = Text() + if max_val <= 0: + max_val = 1.0 + for i in range(0, n, 2): + lv = min(vals[i] / max_val, 1.0) + rv = min(vals[i + 1] / max_val, 1.0) + li = min(int(lv * 4 + 0.5), 4) + ri = min(int(rv * 4 + 0.5), 4) + ch = chr(_BRAILLE_BASE | _LDOTS[li] | _RDOTS[ri]) + result.append(ch, style=_heat(max(lv, rv))) + return result + + +def _fmt_time(seconds: float) -> str: + """Format seconds as MM:SS.s""" + if seconds < 0: + seconds = 0 + m = int(seconds) // 60 + s = seconds - m * 60 + return f"{m:02d}:{s:05.2f}" + + +def _progress_bar(position: float, duration: float, width: int = 40) -> Text: + """Render a progress bar with position indicator.""" + if duration <= 0: + return Text("░" * width, style=theme.DIM) + ratio = min(position / duration, 1.0) + filled = int(ratio * width) + result = Text() + result.append("█" * filled, style=theme.CYAN) + if filled < width: + result.append("▓", style=theme.BRIGHT_CYAN) + result.append("░" * (width - filled - 1), style=theme.DIM) + return result + + +def _short_type(channel: str) -> str: + """Extract the short type name from a channel string.""" + if "#" not in channel: + return "" + return channel.rsplit("#", 1)[-1].rsplit(".", 1)[-1] + + +class RecorderApp(App[None]): + """Terminal VLC for dimos recordings. + + Shows all live LCM topics (like lcmspy). Select topics then press + ``r`` to record, ``space`` to play back, arrow keys to seek, etc. + """ + + CSS_PATH = "../dimos.tcss" + + CSS = f""" + Screen {{ + layout: vertical; + background: {theme.BACKGROUND}; + }} + #streams {{ + height: 1fr; + border: solid {theme.BORDER}; + background: {theme.BG}; + scrollbar-size: 0 0; + }} + #streams > .datatable--header {{ + color: {theme.ACCENT}; + background: transparent; + }} + #streams > .datatable--cursor {{ + background: {theme.BRIGHT_BLACK}; + }} + #timeline {{ + height: 5; + padding: 1 2; + background: {theme.BG}; + border-top: solid {theme.DIM}; + }} + #controls {{ + height: 3; + padding: 0 2; + background: {theme.BG}; + border-top: solid {theme.DIM}; + }} + #status-left {{ + width: 1fr; + }} + #status-right {{ + width: auto; + }} + """ + + BINDINGS = [ + ("q", "quit", "Quit"), + ("space", "toggle_select", "Toggle"), + ("a", "select_all", "All"), + ("n", "select_none", "None"), + ("r", "toggle_record", "Rec"), + ("p", "toggle_play", "Play"), + ("s", "stop_all", "Stop"), + ("left", "seek_back", "-5s"), + ("right", "seek_fwd", "+5s"), + ("shift+left", "seek_back_big", "-30s"), + ("shift+right", "seek_fwd_big", "+30s"), + ("[", "mark_trim_start", "In"), + ("]", "mark_trim_end", "Out"), + ("t", "do_trim", "Trim"), + ("d", "do_delete", "Del"), + ] + + def __init__( + self, + db_path: str | None = None, + autoplay: bool = False, + ) -> None: + super().__init__() + from dimos.protocol.service.lcmservice import autoconf + + autoconf(check_only=True) + + if db_path is None: + ts = time.strftime("%Y%m%d_%H%M%S") + self._db_path = f"recording_{ts}.db" + else: + self._db_path = db_path + self._autoplay = autoplay + self._recorder: RecordReplay | None = None + self._lcm: Any = None + self._spy: Any = None + + # Per-stream sparkline history keyed by stream_name + self._freq_history: dict[str, deque[float]] = {} + # Set of selected stream names (for recording) + self._selected: set[str] = set() + + self._trim_in: float | None = None + self._trim_out: float | None = None + + self._table: DataTable[Any] | None = None + + def compose(self) -> ComposeResult: + yield Header(show_clock=True) + self._table = DataTable(zebra_stripes=False, cursor_type="row") + self._table.id = "streams" + self._table.add_column("REC", key="sel", width=5) + self._table.add_column("Topic", key="topic") + self._table.add_column("Type", key="type") + self._table.add_column("Freq", key="freq") + self._table.add_column("Bandwidth", key="bw") + self._table.add_column("Recorded", key="rec") + self._table.add_column("Activity", key="activity") + yield self._table + yield Static(id="timeline") + with Horizontal(id="controls"): + yield Static(id="status-left") + yield Static(id="status-right") + yield Footer() + + def on_mount(self) -> None: + from dimos.protocol.pubsub.impl.lcmpubsub import LCM + from dimos.record import RecordReplay + from dimos.utils.cli.lcmspy.lcmspy import GraphLCMSpy + + self._lcm = LCM() + + # Live topic discovery via LCM spy (same as lcmspy tool) + self._spy = GraphLCMSpy(graph_log_window=0.5) + self._spy.start() + + if self._db_path: + self._recorder = RecordReplay(self._db_path) + else: + self._recorder = RecordReplay() + + self.title = f"recorder — {self._recorder.path}" + self.set_interval(0.5, self._refresh) + + if self._autoplay and self._db_path: + self._start_playback() + + async def on_unmount(self) -> None: + if self._recorder: + await self._recorder.close() + if self._spy: + self._spy.stop() + if self._lcm and hasattr(self._lcm, "stop"): + self._lcm.stop() + + # ------------------------------------------------------------------ + # Actions + # ------------------------------------------------------------------ + + def action_toggle_select(self) -> None: + """Toggle selection on the row under the cursor.""" + if self._table is None or self._table.row_count == 0: + return + row_key, _ = self._table.coordinate_to_cell_key(self._table.cursor_coordinate) + name = str(row_key.value) + if name in self._selected: + self._selected.discard(name) + else: + self._selected.add(name) + + def action_select_all(self) -> None: + """Select all visible topics.""" + if self._spy is None: + return + with self._spy._topic_lock: + channels = list(self._spy.topic.keys()) + for ch in channels: + self._selected.add(topic_to_stream_name(ch)) + # Also include any already-recorded streams + if self._recorder: + self._selected.update(self._recorder.store.list_streams()) + + def action_select_none(self) -> None: + self._selected.clear() + + def action_toggle_play(self) -> None: + if self._recorder is None: + return + if self._recorder.is_recording: + return + if self._recorder.is_playing: + if self._recorder.is_paused: + self._recorder.resume() + else: + self._recorder.pause() + else: + self._start_playback() + + def action_toggle_record(self) -> None: + if self._recorder is None: + return + if self._recorder.is_playing: + return + if self._recorder.is_recording: + self._recorder.stop_recording() + else: + filt = self._selected if self._selected else None + self._recorder.start_recording([self._lcm], topic_filter=filt) + + async def action_stop_all(self) -> None: + if self._recorder is None: + return + self._recorder.stop_recording() + await self._recorder.stop_playback() + + def action_seek_back(self) -> None: + self._seek_relative(-5.0) + + def action_seek_fwd(self) -> None: + self._seek_relative(5.0) + + def action_seek_back_big(self) -> None: + self._seek_relative(-30.0) + + def action_seek_fwd_big(self) -> None: + self._seek_relative(30.0) + + def action_mark_trim_start(self) -> None: + if self._recorder: + self._trim_in = self._recorder.position + + def action_mark_trim_end(self) -> None: + if self._recorder: + self._trim_out = self._recorder.position + + def action_do_trim(self) -> None: + if self._recorder and self._trim_in is not None and self._trim_out is not None: + self._recorder.stop_playback() + lo, hi = sorted((self._trim_in, self._trim_out)) + self._recorder.trim(lo, hi) + self._trim_in = self._trim_out = None + + def action_do_delete(self) -> None: + if self._recorder and self._trim_in is not None and self._trim_out is not None: + self._recorder.stop_playback() + lo, hi = sorted((self._trim_in, self._trim_out)) + self._recorder.delete_range(lo, hi) + self._trim_in = self._trim_out = None + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _start_playback(self) -> None: + if self._recorder is None or self._lcm is None: + return + if hasattr(self._lcm, "start"): + self._lcm.start() + self._recorder.play(pubsub=self._lcm, speed=1.0) + + def _seek_relative(self, delta: float) -> None: + if self._recorder: + self._recorder.seek(self._recorder.position + delta) + + # ------------------------------------------------------------------ + # Refresh + # ------------------------------------------------------------------ + + def _refresh(self) -> None: + if self._table is None: + return + rec = self._recorder + spy = self._spy + + # ---- Build unified row list: live topics + recorded-only streams ---- + # Each row: (stream_name, channel, spy_topic_or_None) + rows: dict[str, tuple[str, Any]] = {} # stream_name -> (channel, spy_topic) + + if spy is not None: + with spy._topic_lock: + live_topics: dict[str, Any] = dict(spy.topic) # channel -> GraphTopic + for channel, spy_topic in live_topics.items(): + sname = topic_to_stream_name(channel) + rows[sname] = (channel, spy_topic) + + # Add streams that exist in the recording but are not live + recorded_streams = set(rec.store.list_streams()) if rec else set() + for sname in recorded_streams: + if sname not in rows: + rows[sname] = (sname, None) + + # ---- Render table ---- + # Remember cursor position so we can restore it + cursor_row = self._table.cursor_coordinate.row if self._table.row_count > 0 else 0 + self._table.clear(columns=False) + + sorted_names = sorted(rows.keys()) + for sname in sorted_names: + channel, spy_topic = rows[sname] + + # Selection marker + is_sel = sname in self._selected + if is_sel: + sel = Text(" [●] ", style=f"bold {theme.BRIGHT_GREEN}") + else: + sel = Text(" [ ] ", style=theme.DIM) + + # Topic name — green when actively recording, bright when selected + if rec and rec.is_recording and sname in (rec.store.list_streams()): + topic_style = f"bold {theme.BRIGHT_GREEN}" + elif is_sel: + topic_style = theme.BRIGHT_WHITE + else: + topic_style = theme.FOREGROUND + + # Type + type_str = _short_type(channel) if "#" in channel else "" + + # Live freq / bandwidth from spy + if spy_topic is not None: + freq = spy_topic.freq(5.0) + freq_text = Text(f"{freq:.1f} Hz", style=_heat(min(freq / 30.0, 1.0))) + kbps = spy_topic.kbps(5.0) + bw_text = Text(spy_topic.kbps_hr(5.0), style=_heat(min(kbps / 3072, 1.0))) + else: + freq_text = Text("—", style=theme.DIM) + bw_text = Text("—", style=theme.DIM) + + # Recorded count + if sname in recorded_streams: + try: + count = rec.stream(sname).count() + rec_text = Text(str(count), style=theme.YELLOW) + except Exception: + rec_text = Text("—", style=theme.DIM) + else: + rec_text = Text("", style=theme.DIM) + + # Sparkline from spy freq history + if sname not in self._freq_history: + self._freq_history[sname] = deque(maxlen=_SPARK_WIDTH * 2) + if spy_topic is not None: + self._freq_history[sname].append(spy_topic.freq(0.5)) + else: + self._freq_history[sname].append(0.0) + max_f = max(self._freq_history[sname]) if self._freq_history[sname] else 1.0 + activity = _spark(self._freq_history[sname], max_f) + + self._table.add_row( + sel, + Text(sname, style=topic_style), + Text(type_str, style=theme.BLUE), + freq_text, + bw_text, + rec_text, + activity, + key=sname, + ) + + # Restore cursor + if self._table.row_count > 0: + row = min(cursor_row, self._table.row_count - 1) + self._table.move_cursor(row=row) + + # ---- Timeline ---- + duration = rec.duration if rec else 0.0 + position = rec.position if rec else 0.0 + + timeline = Text() + timeline.append(" ") + timeline.append(_fmt_time(position), style=theme.BRIGHT_WHITE) + timeline.append(" ", style=theme.DIM) + timeline.append_text(_progress_bar(position, duration, width=50)) + timeline.append(" ", style=theme.DIM) + timeline.append(_fmt_time(duration), style=theme.FOREGROUND) + + if self._trim_in is not None or self._trim_out is not None: + timeline.append("\n ") + timeline.append("[", style=theme.YELLOW) + timeline.append( + _fmt_time(self._trim_in) if self._trim_in is not None else "--:--", + style=theme.YELLOW if self._trim_in is not None else theme.DIM, + ) + timeline.append(" .. ", style=theme.DIM) + timeline.append( + _fmt_time(self._trim_out) if self._trim_out is not None else "--:--", + style=theme.YELLOW if self._trim_out is not None else theme.DIM, + ) + timeline.append("]", style=theme.YELLOW) + + self.query_one("#timeline", Static).update(timeline) + + # ---- Status bar ---- + status = Text() + if rec and rec.is_recording: + status.append(" ● REC ", style=f"bold on {theme.RED}") + elif rec and rec.is_paused: + status.append(" ❚❚ PAUSED ", style=theme.YELLOW) + elif rec and rec.is_playing: + status.append(" ▶ PLAYING ", style=theme.BRIGHT_GREEN) + else: + status.append(" ■ STOPPED ", style=theme.DIM) + + n_live = len([r for r in rows.values() if r[1] is not None]) + n_sel = len(self._selected) + status.append(f" {n_live} live", style=theme.FOREGROUND) + if n_sel: + status.append(f" {n_sel} selected", style=theme.BRIGHT_GREEN) + if recorded_streams: + status.append(f" {len(recorded_streams)} recorded", style=theme.YELLOW) + + # Contextual hint + if not (rec and (rec.is_recording or rec.is_playing)): + if n_live > 0 and n_sel == 0: + status.append(" SPACE select, A all, R rec", style=theme.DIM) + elif n_sel > 0: + status.append(" R to record selected", style=theme.DIM) + + self.query_one("#status-left", Static).update(status) + + rhs = Text() + if rec: + rhs.append(f"{rec.path} ", style=theme.DIM) + self.query_one("#status-right", Static).update(rhs) + + +def main() -> None: + import sys + + db_path: Path | None = None + autoplay = False + + args = sys.argv[1:] + if args and args[0] == "play" and len(args) > 1: + db_path = Path(args[1]) + autoplay = True + elif args and not args[0].startswith("-"): + db_path = Path(args[0]) + + RecorderApp(db_path=db_path, autoplay=autoplay).run() + + +if __name__ == "__main__": + main() diff --git a/flake.nix b/flake.nix index 68dbf0ee8c..c8c4039e05 100644 --- a/flake.nix +++ b/flake.nix @@ -42,6 +42,7 @@ { vals.pkg=pkgs.opensshWithKerberos;flags={}; } { vals.pkg=pkgs.unixtools.ifconfig; flags={}; } { vals.pkg=pkgs.unixtools.netstat; flags={}; } + { vals.pkg=pkgs.uv; flags={}; } # when pip packages call cc with -I/usr/include, that causes problems on some machines, this swaps that out for the nix cc headers # this is only necessary for pip packages from venv, pip packages from nixpkgs.python312Packages.* already have "-I/usr/include" patched with the nix equivalent diff --git a/pyproject.toml b/pyproject.toml index 1fbd29f86f..3f0833a588 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,6 +87,7 @@ dependencies = [ "psutil>=7.0.0", "sqlite-vec>=0.1.6", "lz4>=4.4.5", + "pytest-asyncio>=0.26.0", ] @@ -99,6 +100,7 @@ dimos = "dimos.robot.cli.dimos:main" rerun-bridge = "dimos.visualization.rerun.bridge:app" doclinks = "dimos.utils.docs.doclinks:main" dtop = "dimos.utils.cli.dtop:main" +recorder = "dimos.utils.cli.recorder.run_recorder:main" [project.urls] Homepage = "https://dimensionalos.com" diff --git a/uv.lock b/uv.lock index 0d6a3a88ab..aba544ae53 100644 --- a/uv.lock +++ b/uv.lock @@ -1700,6 +1700,7 @@ dependencies = [ { name = "psutil" }, { name = "pydantic" }, { name = "pydantic-settings" }, + { name = "pytest-asyncio" }, { name = "python-dotenv" }, { name = "pyturbojpeg" }, { name = "reactivex" }, @@ -2075,6 +2076,7 @@ requires-dist = [ { name = "pymavlink", marker = "extra == 'drone'" }, { name = "pyrealsense2", marker = "sys_platform != 'darwin' and extra == 'manipulation'" }, { name = "pytest", marker = "extra == 'dev'", specifier = "==8.3.5" }, + { name = "pytest-asyncio", specifier = ">=0.26.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = "==0.26.0" }, { name = "pytest-env", marker = "extra == 'dev'", specifier = "==1.1.5" }, { name = "pytest-mock", marker = "extra == 'dev'", specifier = "==3.15.0" }, From 8f763c1a82e6d5c06d6eac98c5f2fd75437c3db8 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 30 Mar 2026 19:47:06 +0100 Subject: [PATCH 02/33] Fix async --- dimos/record/record_replay.py | 4 ++-- dimos/record/test_record_replay.py | 2 +- dimos/utils/cli/recorder/run_recorder.py | 28 ++++++++++++------------ 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index 4ef31e7511..888f28fd1a 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -331,7 +331,7 @@ async def stop_playback(self) -> None: await self._play_task self._play_task = None - def seek(self, seconds: float) -> None: + async def seek(self, seconds: float) -> None: """Set playback position in seconds from recording start. Takes effect immediately if playing (restarts playback loop). @@ -339,7 +339,7 @@ def seek(self, seconds: float) -> None: """ self._position = max(0.0, min(seconds, self.duration)) if self.is_playing: - self.stop_playback() + await self.stop_playback() assert self._pubsub is not None self.play(pubsub=self._pubsub, speed=self._play_speed) diff --git a/dimos/record/test_record_replay.py b/dimos/record/test_record_replay.py index a14d1ef3cd..f75ef121a2 100644 --- a/dimos/record/test_record_replay.py +++ b/dimos/record/test_record_replay.py @@ -238,7 +238,7 @@ async def test_seek(self, tmp_db: str) -> None: await asyncio.sleep(0.01) rec.stop_recording() - rec.seek(0.05) + await rec.seek(0.05) assert rec.position == pytest.approx(0.05, abs=0.01) async def test_multiple_pubsubs(self, tmp_db: str) -> None: diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index 4827abf6f2..4f47cf1ce3 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -311,17 +311,17 @@ async def action_stop_all(self) -> None: self._recorder.stop_recording() await self._recorder.stop_playback() - def action_seek_back(self) -> None: - self._seek_relative(-5.0) + async def action_seek_back(self) -> None: + await self._seek_relative(-5.0) - def action_seek_fwd(self) -> None: - self._seek_relative(5.0) + async def action_seek_fwd(self) -> None: + await self._seek_relative(5.0) - def action_seek_back_big(self) -> None: - self._seek_relative(-30.0) + async def action_seek_back_big(self) -> None: + await self._seek_relative(-30.0) - def action_seek_fwd_big(self) -> None: - self._seek_relative(30.0) + async def action_seek_fwd_big(self) -> None: + await self._seek_relative(30.0) def action_mark_trim_start(self) -> None: if self._recorder: @@ -331,16 +331,16 @@ def action_mark_trim_end(self) -> None: if self._recorder: self._trim_out = self._recorder.position - def action_do_trim(self) -> None: + async def action_do_trim(self) -> None: if self._recorder and self._trim_in is not None and self._trim_out is not None: - self._recorder.stop_playback() + await self._recorder.stop_playback() lo, hi = sorted((self._trim_in, self._trim_out)) self._recorder.trim(lo, hi) self._trim_in = self._trim_out = None - def action_do_delete(self) -> None: + async def action_do_delete(self) -> None: if self._recorder and self._trim_in is not None and self._trim_out is not None: - self._recorder.stop_playback() + await self._recorder.stop_playback() lo, hi = sorted((self._trim_in, self._trim_out)) self._recorder.delete_range(lo, hi) self._trim_in = self._trim_out = None @@ -356,9 +356,9 @@ def _start_playback(self) -> None: self._lcm.start() self._recorder.play(pubsub=self._lcm, speed=1.0) - def _seek_relative(self, delta: float) -> None: + async def _seek_relative(self, delta: float) -> None: if self._recorder: - self._recorder.seek(self._recorder.position + delta) + await self._recorder.seek(self._recorder.position + delta) # ------------------------------------------------------------------ # Refresh From 4e646d6d7537e96786ecd9ea3e6fa2c446875f2d Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Tue, 31 Mar 2026 19:52:01 +0100 Subject: [PATCH 03/33] Load in viewer as separate source --- dimos/record/record_replay.py | 90 +++++++++++++++-------------- dimos/visualization/rerun/bridge.py | 2 +- 2 files changed, 47 insertions(+), 45 deletions(-) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index 888f28fd1a..d769f039c8 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -47,10 +47,11 @@ import time from typing import Any, TypedDict -from dimos.memory2.codecs.base import _resolve_payload_type +import rerun as rr + from dimos.memory2.store.sqlite import SqliteStore from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic -from dimos.protocol.pubsub.spec import PubSub +from dimos.visualization.rerun.bridge import RerunConvertible, is_rerun_multi if sys.version_info >= (3, 11): from typing import Self @@ -164,17 +165,21 @@ def stop_recording(self) -> None: logger.info("Recording stopped") def _on_message(self, msg: bytes, topic: Topic) -> None: - """Handle incoming message during recording.""" stream_name = topic_to_stream_name(topic.pattern) if self._topic_filter is not None and stream_name not in self._topic_filter: return - msg_type = type(msg) - - s = self._store.stream(stream_name, msg_type) + s = self._store.stream(stream_name, type(msg)) s.append(msg, ts=time.time()) + # Persist the full channel string (with #type) in the registry + # so playback can reconstruct the lcm_type for decoding. + reg = self._store._registry.get(stream_name) + if reg and "channel" not in reg: + reg["channel"] = str(topic) + self._store._registry.put(stream_name, reg) + @property def duration(self) -> float: """Total duration of the recording in seconds.""" @@ -222,34 +227,40 @@ def stream_info(self) -> tuple[StreamInfo, ...]: result.append(info) return tuple(result) - def play( - self, - pubsub: PubSub[Any, Any] | None = None, - speed: float = 1.0, - ) -> None: - """Start playback in a background thread. + def play(self, speed: float = 1.0) -> None: + """Start playback as a separate Rerun recording. - Args: - pubsub: If provided, messages are published to this pubsub. - Use LCM() to make them visible to rerun-bridge. - speed: Playback speed multiplier (1.0 = realtime, 2.0 = 2x, etc). + Connects to the running Rerun viewer and logs decoded messages + under a recording called ``"playback"``, so it appears alongside + (but separate from) any live data. """ if self.is_playing: raise RuntimeError("Already playing") self._play_speed = speed - self._pubsub = pubsub # Set resume so playback starts, this is cleared to pause playback. self._resume.set() - self._play_task = asyncio.create_task(self._playback_loop(pubsub)) + self._play_task = asyncio.create_task(self._playback_loop()) - async def _playback_loop(self, pubsub: PubSub[Any, Any] | None) -> None: - """Main playback loop — merges all streams by timestamp.""" + async def _playback_loop(self) -> None: t_min, t_max = self.time_range if t_min >= t_max: return - # Build iterators for each stream, starting from seek position + # Separate Rerun recording so playback appears as its own source + rec = rr.RecordingStream("playback", make_default=False) + rec.connect_grpc() + + # Build topic map for decoding raw bytes -> DimosMsg + topic_map: dict[str, Topic] = {} + for name in self._store.list_streams(): + reg = self._store._registry.get(name) + if reg: + channel = reg.get("channel") + if channel: + topic_map[name] = Topic.from_channel_str(channel) + + # Merge-sort all streams by timestamp start_ts = t_min + self._position heap: list[tuple[float, int, str, Any]] = [] counter = 0 # tiebreaker for heapq @@ -267,7 +278,6 @@ async def _playback_loop(self, pubsub: PubSub[Any, Any] | None) -> None: if not heap: return - topic_map = {} if pubsub is None else self._build_topic_map(pubsub) wall_start = time.monotonic() rec_start = heap[0][0] # earliest observation timestamp @@ -282,12 +292,22 @@ async def _playback_loop(self, pubsub: PubSub[Any, Any] | None) -> None: # Wait for correct playback time elapsed_rec = ts - rec_start target_wall = wall_start + (elapsed_rec / self._play_speed) - sleep_time = target_wall - time.monotonic() - await asyncio.sleep(sleep_time) + await asyncio.sleep(target_wall - time.monotonic()) self._position = ts - t_min - if pubsub is not None and stream_name in topic_map: - pubsub.publish(topic_map[stream_name], obs.data) + + # Decode raw bytes -> DimosMsg -> Rerun archetype + topic = topic_map.get(stream_name) + if topic is not None and topic.lcm_type is not None: + msg = topic.lcm_type.lcm_decode(obs.data) + if isinstance(msg, RerunConvertible): + entity_path = f"world/{stream_name}" + rerun_data = msg.to_rerun() + if is_rerun_multi(rerun_data): + for path, archetype in rerun_data: + rec.log(path, archetype) + else: + rec.log(entity_path, rerun_data) try: next_obs = next(it) @@ -296,24 +316,6 @@ async def _playback_loop(self, pubsub: PubSub[Any, Any] | None) -> None: heapq.heappush(heap, (next_obs.ts, counter, stream_name, (next_obs, it))) counter += 1 - def _build_topic_map(self, pubsub: PubSub[Any, Any]) -> dict[str, Topic]: - """Build stream_name -> Topic mapping for publishing.""" - topic_map: dict[str, Topic] = {} - - for name in self._store.list_streams(): - reg = self._store._registry.get(name) - if reg is None: - continue - - payload_module = reg.get("payload_module") - lcm_type = None - if payload_module: - lcm_type = _resolve_payload_type(payload_module) - - topic_map[name] = Topic(f"/{name}", lcm_type=lcm_type) - - return topic_map - def pause(self) -> None: """Pause playback. Resume with :meth:`resume`.""" self._resume.clear() diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 8b1cda443c..f793b9966c 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -307,7 +307,7 @@ def start(self) -> None: ) rr.spawn(connect=True, memory_limit=self.config.memory_limit) elif self.config.viewer_mode == "web": - server_uri = rr.serve_grpc() + server_uri = rr.serve_grpc(grpc_port=RERUN_GRPC_PORT) rr.serve_web_viewer(connect_to=server_uri, open_browser=False) elif self.config.viewer_mode == "connect": rr.connect_grpc(self.config.connect_url) From 6c7ff5dc48295aa93826bef0bbd77536a886b300 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Wed, 1 Apr 2026 00:30:38 +0100 Subject: [PATCH 04/33] Fix playback --- dimos/utils/cli/recorder/run_recorder.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index 4f47cf1ce3..a1f84d4bbe 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -222,11 +222,11 @@ def compose(self) -> ComposeResult: yield Footer() def on_mount(self) -> None: - from dimos.protocol.pubsub.impl.lcmpubsub import LCM + from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase from dimos.record import RecordReplay from dimos.utils.cli.lcmspy.lcmspy import GraphLCMSpy - self._lcm = LCM() + self._lcm = LCMPubSubBase() # Live topic discovery via LCM spy (same as lcmspy tool) self._spy = GraphLCMSpy(graph_log_window=0.5) @@ -354,7 +354,7 @@ def _start_playback(self) -> None: return if hasattr(self._lcm, "start"): self._lcm.start() - self._recorder.play(pubsub=self._lcm, speed=1.0) + self._recorder.play(speed=1.0) async def _seek_relative(self, delta: float) -> None: if self._recorder: From ff0adfcd393b5fb63256d42138b72146c8ac4653 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sat, 4 Apr 2026 13:43:00 +0100 Subject: [PATCH 05/33] Use AllPubSub --- dimos/record/record_replay.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index d769f039c8..d9c8bf9361 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -50,8 +50,9 @@ import rerun as rr from dimos.memory2.store.sqlite import SqliteStore -from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase, Topic -from dimos.visualization.rerun.bridge import RerunConvertible, is_rerun_multi +from dimos.protocol.pubsub.impl.lcmpubsub import Topic +from dimos.protocol.pubsub.spec import AllPubSub +from dimos.protocol.service.spec import Service if sys.version_info >= (3, 11): from typing import Self @@ -129,18 +130,12 @@ def is_paused(self) -> bool: def start_recording( self, - pubsubs: Collection[LCMPubSubBase], + pubsubs: Collection[AllPubSub[Any, Any]], topic_filter: Container[str] | None = None, ) -> None: """Start recording messages from the given pubsubs. - Each pubsub is subscribed via ``subscribe_all()``. Messages are - stored in per-topic streams with automatic codec selection. - - Args: - pubsubs: List of pubsubs to subscribe to. - topic_filter: If provided, only record topics whose sanitized - stream name is in this set. If ``None``, record everything. + Accepts any ``AllPubSub`` (LCM, SHM, etc). """ if self._recording: raise RuntimeError("Already recording") @@ -148,7 +143,8 @@ def start_recording( self._topic_filter = topic_filter for pubsub in pubsubs: - pubsub.start() + if isinstance(pubsub, Service): + pubsub.start() unsub = pubsub.subscribe_all(self._on_message) self._unsubscribes.append(unsub) From 90830968dd6fea6c0ed96cff29404c250c55afdc Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 6 Apr 2026 16:45:28 +0100 Subject: [PATCH 06/33] Recording fixes --- dimos/core/blueprints.py | 6 +-- dimos/memory2/store/base.py | 8 +-- dimos/protocol/pubsub/spec.py | 7 ++- dimos/record/record_replay.py | 62 +++++++++++++----------- dimos/utils/cli/recorder/run_recorder.py | 31 ++++++++++-- 5 files changed, 70 insertions(+), 44 deletions(-) diff --git a/dimos/core/blueprints.py b/dimos/core/blueprints.py index 55e152109e..79904e94fb 100644 --- a/dimos/core/blueprints.py +++ b/dimos/core/blueprints.py @@ -30,7 +30,6 @@ from dimos.core.module_coordinator import ModuleCoordinator from dimos.core.stream import In, Out from dimos.core.transport import LCMTransport, PubSubTransport, pLCMTransport -from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.record.record_replay import RecordReplay from dimos.spec.utils import Spec, is_spec, spec_annotation_compliance, spec_structural_compliance from dimos.utils.generic import short_id @@ -510,10 +509,7 @@ def replay( patched = self.disabled_modules(*modules_to_disable) coordinator = patched.build(cli_config_overrides) - # Start playback in background — publishes to LCM so other modules receive data - lcm = LCM() - lcm.start() - recording.play(pubsub=lcm, speed=speed) + recording.play(speed=speed) return coordinator diff --git a/dimos/memory2/store/base.py b/dimos/memory2/store/base.py index cf571f23b0..dad75d5636 100644 --- a/dimos/memory2/store/base.py +++ b/dimos/memory2/store/base.py @@ -14,7 +14,7 @@ from __future__ import annotations -from typing import Any, TypeVar, cast +from typing import Any, TypeVar from dimos.core.resource import CompositeResource from dimos.memory2.backend import Backend @@ -97,7 +97,7 @@ def streams(self) -> StreamAccessor: @staticmethod def _resolve_codec( - payload_type: type[Any] | None, raw_codec: Codec[Any] | str | None + payload_type: type[object] | None, raw_codec: Codec[object] | str | None ) -> Codec[Any]: if isinstance(raw_codec, Codec): return raw_codec @@ -111,7 +111,7 @@ def _resolve_codec( return codec_for(payload_type) def _create_backend( - self, name: str, payload_type: type[Any] | None = None, **config: Any + self, name: str, payload_type: type[object] | None = None, **config: Any ) -> Backend[Any]: """Create a Backend for the named stream. Called once per stream name.""" codec = self._resolve_codec(payload_type, config.pop("codec", None)) @@ -155,7 +155,7 @@ def stream(self, name: str, payload_type: type[T] | None = None, **overrides: An resolved = {**self.config.model_dump(exclude_none=True), **overrides} backend = self._create_backend(name, payload_type, **resolved) self._streams[name] = Stream(source=backend) - return cast("Stream[T]", self._streams[name]) + return self._streams[name] def list_streams(self) -> list[str]: """Return names of all streams in this store.""" diff --git a/dimos/protocol/pubsub/spec.py b/dimos/protocol/pubsub/spec.py index 0e292cfcb7..63709b39e0 100644 --- a/dimos/protocol/pubsub/spec.py +++ b/dimos/protocol/pubsub/spec.py @@ -17,10 +17,13 @@ from collections.abc import AsyncIterator, Callable from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, Generic, Protocol, TypeVar, runtime_checkable +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable + +if TYPE_CHECKING: + from dimos.protocol.pubsub.impl.lcmpubsub import Topic MsgT = TypeVar("MsgT") -TopicT = TypeVar("TopicT") +TopicT = TypeVar("TopicT", bound="Topic") MsgT_co = TypeVar("MsgT_co", covariant=True) TopicT_co = TypeVar("TopicT_co", covariant=True) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index d9c8bf9361..217007f528 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -37,7 +37,7 @@ """ import asyncio -from collections.abc import Callable, Collection, Container +from collections.abc import Callable, Collection from contextlib import suppress import heapq import logging @@ -45,14 +45,17 @@ import re import sys import time -from typing import Any, TypedDict +from typing import Any, NotRequired, TypedDict import rerun as rr from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.stream import Stream from dimos.protocol.pubsub.impl.lcmpubsub import Topic -from dimos.protocol.pubsub.spec import AllPubSub +from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.service.spec import Service +from dimos.types.timestamped import Timestamped +from dimos.visualization.rerun.bridge import RerunConvertible, is_rerun_multi if sys.version_info >= (3, 11): from typing import Self @@ -76,10 +79,10 @@ def topic_to_stream_name(channel: str) -> str: class StreamInfo(TypedDict): name: str count: int - start: float - end: float - duration: float - type: str + start: NotRequired[float] + end: NotRequired[float] + duration: NotRequired[float] + type: NotRequired[str] class RecordReplay: @@ -98,10 +101,10 @@ def __init__(self, path: str) -> None: self._recording = False self._unsubscribes: list[Callable[[], None]] = [] - self._topic_filter: Container[str] | None = None + self._ts_offsets: dict[str, float] = {} self._resume = asyncio.Event() - self._play_task: asyncio.Task | None = None + self._play_task: asyncio.Task[None] | None = None self._play_speed = 1.0 self._position = 0.0 self._pubsub = None @@ -130,47 +133,48 @@ def is_paused(self) -> bool: def start_recording( self, - pubsubs: Collection[AllPubSub[Any, Any]], - topic_filter: Container[str] | None = None, + pubsubs: Collection[PubSub[Topic, object]], + topics: Collection[Topic], ) -> None: - """Start recording messages from the given pubsubs. - - Accepts any ``AllPubSub`` (LCM, SHM, etc). - """ + """Start recording messages from the given pubsubs.""" if self._recording: raise RuntimeError("Already recording") self._recording = True - self._topic_filter = topic_filter for pubsub in pubsubs: if isinstance(pubsub, Service): pubsub.start() - unsub = pubsub.subscribe_all(self._on_message) - self._unsubscribes.append(unsub) + for topic in topics: + unsub = pubsub.subscribe(topic, self._on_message) + self._unsubscribes.append(unsub) logger.info("Recording started on %d pubsub(s)", len(pubsubs)) def stop_recording(self) -> None: - """Stop recording.""" if not self._recording: return self._recording = False for unsub in self._unsubscribes: unsub() self._unsubscribes.clear() + self._ts_offsets.clear() logger.info("Recording stopped") - def _on_message(self, msg: bytes, topic: Topic) -> None: + def _on_message(self, msg: object, topic: Topic) -> None: stream_name = topic_to_stream_name(topic.pattern) - if self._topic_filter is not None and stream_name not in self._topic_filter: - return + ts: float | None = None + if isinstance(msg, Timestamped) and msg.ts: + # Record the time offset on first message and then offset all messages + # by the same amount. This allows recording live channels and replayed + # channels together with consistent timestamps. + if stream_name not in self._ts_offsets: + self._ts_offsets[stream_name] = time.time() - msg.ts + ts = msg.ts + self._ts_offsets[stream_name] s = self._store.stream(stream_name, type(msg)) - s.append(msg, ts=time.time()) + s.append(msg, ts=ts) - # Persist the full channel string (with #type) in the registry - # so playback can reconstruct the lcm_type for decoding. reg = self._store._registry.get(stream_name) if reg and "channel" not in reg: reg["channel"] = str(topic) @@ -191,7 +195,7 @@ def time_range(self) -> tuple[float, float]: t_min = math.inf t_max = -math.inf for name in streams: - s = self._store.stream(name) + s: Stream[object] = self._store.stream(name) if s.exists(): t0, t1 = s.get_time_range() t_min = min(t_min, t0) @@ -209,7 +213,7 @@ def stream_info(self) -> tuple[StreamInfo, ...]: """Return per-stream metadata: name, count, time range, type.""" result = [] for name in self._store.list_streams(): - s = self._store.stream(name) + s: Stream[object] = self._store.stream(name) info: StreamInfo = {"name": name, "count": s.count()} if info["count"] > 0: t0, t1 = s.get_time_range() @@ -262,7 +266,7 @@ async def _playback_loop(self) -> None: counter = 0 # tiebreaker for heapq for name in self._store.list_streams(): - s = self._store.stream(name) + s: Stream[object] = self._store.stream(name) it = iter(s.after(start_ts - 0.001)) try: obs = next(it) @@ -352,7 +356,7 @@ def delete_range(self, start: float, end: float) -> int: total = 0 for name in self._store.list_streams(): - s = self._store.stream(name) + s: Stream[object] = self._store.stream(name) total += s.delete_range(abs_start, abs_end) return total diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index a1f84d4bbe..f743b2fc62 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -27,6 +27,7 @@ from __future__ import annotations from collections import deque +from pathlib import Path import time from typing import Any @@ -36,6 +37,7 @@ from textual.containers import Horizontal from textual.widgets import DataTable, Footer, Header, Static +from dimos.record.record_replay import RecordReplay from dimos.utils.cli import theme # Braille sparkline constants (same as dtop) @@ -222,11 +224,11 @@ def compose(self) -> ComposeResult: yield Footer() def on_mount(self) -> None: - from dimos.protocol.pubsub.impl.lcmpubsub import LCMPubSubBase + from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.record import RecordReplay from dimos.utils.cli.lcmspy.lcmspy import GraphLCMSpy - self._lcm = LCMPubSubBase() + self._lcm = LCM() # Live topic discovery via LCM spy (same as lcmspy tool) self._spy = GraphLCMSpy(graph_log_window=0.5) @@ -302,8 +304,8 @@ def action_toggle_record(self) -> None: if self._recorder.is_recording: self._recorder.stop_recording() else: - filt = self._selected if self._selected else None - self._recorder.start_recording([self._lcm], topic_filter=filt) + topics = self._selected_topics() if self._selected else self._all_topics() + self._recorder.start_recording([self._lcm], topics=topics) async def action_stop_all(self) -> None: if self._recorder is None: @@ -349,6 +351,27 @@ async def action_do_delete(self) -> None: # Helpers # ------------------------------------------------------------------ + def _all_topics(self) -> list[Any]: + """All discovered topics from the spy.""" + from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic + + if not self._spy: + return [] + return [LCMTopic.from_channel_str(ch) for ch in self._spy.topic] + + def _selected_topics(self) -> list[Any]: + """Map selected stream names back to LCM Topics via the spy.""" + from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic + from dimos.record.record_replay import topic_to_stream_name + + if not self._spy: + return [] + return [ + LCMTopic.from_channel_str(ch) + for ch in self._spy.topic + if topic_to_stream_name(ch) in self._selected + ] + def _start_playback(self) -> None: if self._recorder is None or self._lcm is None: return From 831517844c570294535ac1cd4782981fee26f0e1 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Tue, 7 Apr 2026 15:22:16 +0100 Subject: [PATCH 07/33] Fix playback --- dimos/record/record_replay.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index 217007f528..025f84aafd 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -296,10 +296,13 @@ async def _playback_loop(self) -> None: self._position = ts - t_min - # Decode raw bytes -> DimosMsg -> Rerun archetype topic = topic_map.get(stream_name) if topic is not None and topic.lcm_type is not None: - msg = topic.lcm_type.lcm_decode(obs.data) + msg = ( + obs.data + if not isinstance(obs.data, bytes) + else topic.lcm_type.lcm_decode(obs.data) + ) if isinstance(msg, RerunConvertible): entity_path = f"world/{stream_name}" rerun_data = msg.to_rerun() From 46b38b89321e35b740fc460cba13c6aacc14faac Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 10 Apr 2026 18:00:40 +0100 Subject: [PATCH 08/33] Fix duplicate recording timestamps --- dimos/experimental/security_demo/security_module.py | 2 ++ .../frontier_exploration/wavefront_frontier_goal_selector.py | 2 -- dimos/navigation/patrolling/module.py | 3 +++ dimos/robot/unitree/go2/connection.py | 2 ++ 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/dimos/experimental/security_demo/security_module.py b/dimos/experimental/security_demo/security_module.py index dc0c29b9da..a07f9794dc 100644 --- a/dimos/experimental/security_demo/security_module.py +++ b/dimos/experimental/security_demo/security_module.py @@ -372,6 +372,8 @@ def _cancel_current_goal(self) -> None: with self._lock: pose = self._latest_pose if pose is not None: + # We want to update timestamps, otherwise recordings etc. would fail + pose.ts = time.time() self.goal_request.publish(pose) def _transition_to(self, new_state: State) -> None: diff --git a/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py b/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py index b8dbe0dfc8..a1a62c76f0 100644 --- a/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py +++ b/dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py @@ -741,7 +741,6 @@ def stop_exploration(self) -> bool: position=self.latest_odometry.position, orientation=self.latest_odometry.orientation, frame_id="world", - ts=self.latest_odometry.ts, ) self.goal_request.publish(goal) @@ -782,7 +781,6 @@ def _exploration_loop(self) -> None: goal_msg.position.z = 0.0 goal_msg.orientation.w = 1.0 # No rotation goal_msg.frame_id = "world" - goal_msg.ts = self.latest_costmap.ts self.goal_request.publish(goal_msg) logger.info(f"Published frontier goal: ({goal.x:.2f}, {goal.y:.2f})") diff --git a/dimos/navigation/patrolling/module.py b/dimos/navigation/patrolling/module.py index 72a2b821fc..a46fc1fe40 100644 --- a/dimos/navigation/patrolling/module.py +++ b/dimos/navigation/patrolling/module.py @@ -14,6 +14,7 @@ import threading +import time from dimos_lcm.std_msgs import Bool from reactivex.disposable import Disposable @@ -138,6 +139,8 @@ def _stop_patrolling(self) -> None: # Publish current position as goal to cancel in-progress navigation. pose = self._latest_pose if pose is not None: + # We want to update timestamps, otherwise recordings etc. would fail + pose.ts = time.time() self.goal_request.publish(pose) with self._patrol_lock: if self._patrol_thread is not None: diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index b994f0cae0..afa153369b 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -299,6 +299,8 @@ def _publish_tf(self, msg: PoseStamped) -> None: def publish_camera_info(self) -> None: while True: + # We want to update timestamps, otherwise recordings etc. would fail + self.camera_info_static.ts = time.time() self.camera_info.publish(self.camera_info_static) time.sleep(1.0) From 99e6013dbf40dc75395c22eda2be660a3dcaa741 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 10 Apr 2026 21:27:20 +0100 Subject: [PATCH 09/33] Fixes --- dimos/core/coordination/blueprints.py | 5 +++++ dimos/core/coordination/module_coordinator.py | 11 ++++++++--- dimos/core/module.py | 3 +-- .../unitree/go2/blueprints/basic/unitree_go2_basic.py | 1 + 4 files changed, 15 insertions(+), 5 deletions(-) diff --git a/dimos/core/coordination/blueprints.py b/dimos/core/coordination/blueprints.py index dd7686c266..128b94688e 100644 --- a/dimos/core/coordination/blueprints.py +++ b/dimos/core/coordination/blueprints.py @@ -151,6 +151,7 @@ class Blueprint: ) requirement_checks: tuple[Callable[[], str | None], ...] = field(default_factory=tuple) configurator_checks: "tuple[SystemConfigurator, ...]" = field(default_factory=tuple) + record_modules: tuple[type[ModuleBase], ...] = field(default_factory=tuple) @classmethod def create(cls, module: type[ModuleBase], **kwargs: Any) -> "Blueprint": @@ -160,6 +161,9 @@ def create(cls, module: type[ModuleBase], **kwargs: Any) -> "Blueprint": def disabled_modules(self, *modules: type[ModuleBase]) -> "Blueprint": return replace(self, disabled_modules_tuple=self.disabled_modules_tuple + modules) + def default_record_modules(self, *modules: type[ModuleBase]) -> "Blueprint": + return replace(self, record_modules=self.record_modules + modules) + def transports(self, transports: dict[tuple[str, type], Any]) -> "Blueprint": return replace(self, transport_map=MappingProxyType({**self.transport_map, **transports})) @@ -216,6 +220,7 @@ def autoconnect(*blueprints: Blueprint) -> Blueprint: remapping_map=MappingProxyType(all_remappings), requirement_checks=all_requirement_checks, configurator_checks=all_configurator_checks, + record_modules=tuple(module for bp in blueprints for module in bp.record_modules), ) diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index fbd7dff81f..d0d3f25426 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -19,6 +19,7 @@ from collections.abc import Mapping import dataclasses import importlib +from pathlib import Path import shutil import sys from types import MappingProxyType @@ -271,10 +272,14 @@ def build( coordinator.start_all_modules() if global_config.record_path: + # Delete existing file, don't append to it. + Path(global_config.record_path).unlink(missing_ok=True) + record_modules = blueprint.record_modules for bp in blueprint.active_blueprints: - instance = coordinator.get_instance(bp.module) - if instance is not None: - instance.start_recording(global_config.record_path) + if bp.module in record_modules: + instance = coordinator.get_instance(bp.module) + if instance is not None: + instance.start_recording(global_config.record_path) _log_blueprint_graph(blueprint, coordinator) diff --git a/dimos/core/module.py b/dimos/core/module.py index 40a7c4d8ce..1b3efa7c0f 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -389,10 +389,9 @@ def _start_recording(self, db_path: str) -> None: self._rec_unsubs: list[Callable[[], None]] = [] for name, out in self.outputs.items(): stream = self._rec_store.stream(name, out.type) - payload_mod = f"{out.type.__module__}.{out.type.__qualname__}" reg = self._rec_store._registry.get(name) if reg and "channel" not in reg: - reg["channel"] = f"/{name}#{payload_mod}" + reg["channel"] = f"/{name}#{out.type.msg_name}" self._rec_store._registry.put(name, reg) def cb(msg: Any, _stream: Any = stream) -> None: 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 ab55b7dbb6..8ba0b45600 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -147,6 +147,7 @@ def _go2_rerun_blueprint() -> Any: ) .global_config(n_workers=4, robot_model="unitree_go2") .configurators(ClockSyncConfigurator()) + .default_record_modules(GO2Connection) ) __all__ = [ From 81157b4510f07ff9b8a2c41ce9992937a3dfd8d5 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 10 Apr 2026 21:29:57 +0100 Subject: [PATCH 10/33] Revert flake.nix --- flake.nix | 1 - 1 file changed, 1 deletion(-) diff --git a/flake.nix b/flake.nix index 89c7190349..c22b1f7791 100644 --- a/flake.nix +++ b/flake.nix @@ -42,7 +42,6 @@ { vals.pkg=pkgs.opensshWithKerberos;flags={}; } { vals.pkg=pkgs.unixtools.ifconfig; flags={}; } { vals.pkg=pkgs.unixtools.netstat; flags={}; } - { vals.pkg=pkgs.uv; flags={}; } # when pip packages call cc with -I/usr/include, that causes problems on some machines, this swaps that out for the nix cc headers # this is only necessary for pip packages from venv, pip packages from nixpkgs.python312Packages.* already have "-I/usr/include" patched with the nix equivalent From 9d6957612cb7ea4d933969340ddd5066142aecd0 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 10 Apr 2026 22:08:23 +0100 Subject: [PATCH 11/33] Cleanup --- dimos/core/coordination/module_coordinator.py | 2 +- dimos/core/module.py | 46 ++++++------------- 2 files changed, 16 insertions(+), 32 deletions(-) diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index d0d3f25426..bfcc7fbf5c 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -321,7 +321,7 @@ def replay( if not modules_to_disable: logger.warning( - "Replay: no modules disables - recording streams %s" + "Replay: no modules disabled - recording streams %s" "don't match any module OUT names", recorded_streams, ) diff --git a/dimos/core/module.py b/dimos/core/module.py index 1b3efa7c0f..040d4aae9d 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -19,6 +19,7 @@ import json import sys import threading +import time from typing import ( TYPE_CHECKING, Any, @@ -39,10 +40,13 @@ from dimos.core.resource import CompositeResource from dimos.core.rpc_client import RpcCall from dimos.core.stream import In, Out, RemoteOut, Transport +from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.stream import Stream from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable from dimos.protocol.tf.tf import LCMTF, TFSpec +from dimos.types.timestamped import Timestamped from dimos.utils import colors from dimos.utils.generic import classproperty @@ -112,9 +116,11 @@ class ModuleBase(Configurable, CompositeResource): _module_closed: bool = False _module_closed_lock: threading.Lock _loop_thread_timeout: float = 2.0 + _rec_store: SqliteStore | None = None def __init__(self, config_args: dict[str, Any]) -> None: super().__init__(**config_args) + self._rec_unsubs: list[Callable[[], None]] = [] self._module_closed_lock = threading.Lock() self._loop, self._loop_thread = get_loop() try: @@ -145,7 +151,7 @@ def build(self) -> None: @rpc def start(self) -> None: - self._listen_for_recording_commands() + pass @rpc def stop(self) -> None: @@ -374,19 +380,7 @@ def blueprint(self) -> _BlueprintPartial: @rpc def start_recording(self, db_path: str) -> None: - self._start_recording(db_path) - - @rpc - def stop_recording(self) -> None: - self._stop_recording() - - def _start_recording(self, db_path: str) -> None: - import time - - from dimos.memory2.store.sqlite import SqliteStore - self._rec_store = SqliteStore(path=db_path) - self._rec_unsubs: list[Callable[[], None]] = [] for name, out in self.outputs.items(): stream = self._rec_store.stream(name, out.type) reg = self._rec_store._registry.get(name) @@ -394,34 +388,24 @@ def _start_recording(self, db_path: str) -> None: reg["channel"] = f"/{name}#{out.type.msg_name}" self._rec_store._registry.put(name, reg) - def cb(msg: Any, _stream: Any = stream) -> None: - ts = getattr(msg, "ts", None) or time.time() + def cb(msg: Any, _stream: Stream[object] = stream) -> None: + ts = msg.ts if isinstance(msg, Timestamped) else time.time() _stream.append(msg, ts=ts) self._rec_unsubs.append(out.subscribe(cb)) + @rpc + def stop_recording(self) -> None: + self._stop_recording() + def _stop_recording(self) -> None: - for unsub in getattr(self, "_rec_unsubs", []): + for unsub in self._rec_unsubs: unsub() self._rec_unsubs = [] - if hasattr(self, "_rec_store") and self._rec_store: + if self._rec_store is not None: self._rec_store.stop() self._rec_store = None - def _listen_for_recording_commands(self) -> None: - from dimos.protocol.pubsub.impl.lcmpubsub import PickleLCM, Topic as LCMTopic - - self._rec_cmd_lcm = PickleLCM() - self._rec_cmd_lcm.start() - - def on_cmd(msg: dict, topic: Any) -> None: - if msg.get("action") == "start": - self._start_recording(msg["db_path"]) - elif msg.get("action") == "stop": - self._stop_recording() - - self._rec_cmd_lcm.subscribe(LCMTopic("/dimos/record_cmd"), on_cmd) - @rpc def set_module_ref(self, name: str, module_ref: "RPCClient") -> None: setattr(self, name, module_ref) From 6bf01dadba68909319d20065133d95981d608f10 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 10 Apr 2026 22:30:06 +0100 Subject: [PATCH 12/33] Import time --- dimos/visualization/rerun/bridge.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 2f07c8806e..dd2dbfc361 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -41,7 +41,6 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig -from dimos.msgs.sensor_msgs.PointCloud2 import register_colormap_annotation from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.pubsub.patterns import Glob, pattern_matches from dimos.protocol.pubsub.spec import SubscribeAllCapable @@ -290,6 +289,9 @@ def _on_message(self, msg: Any, topic: Any) -> None: @rpc def start(self) -> None: + # Delay import to reduce import time (~2.4s) + from dimos.msgs.sensor_msgs.PointCloud2 import register_colormap_annotation + super().start() logger.info("Rerun bridge starting", viewer_mode=self.config.viewer_mode) From 7b6d0c1105f65afc2c2bfe8c59b21371be1d7184 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 10 Apr 2026 23:36:39 +0100 Subject: [PATCH 13/33] Docs --- docs/usage/blueprints.md | 36 ++++++++++++++++++++++++++++++++---- docs/usage/cli.md | 19 ++++++++++++++++++- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/docs/usage/blueprints.md b/docs/usage/blueprints.md index f6bbf60779..a75881bfbc 100644 --- a/docs/usage/blueprints.md +++ b/docs/usage/blueprints.md @@ -7,7 +7,7 @@ You don't typically want to run a single module, so multiple blueprints are hand You create a `Blueprint` from a single module (say `ConnectionModule`) with: ```python session=blueprint-ex1 -from dimos.core.blueprints import Blueprint +from dimos.core.coordination.blueprints import Blueprint from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig @@ -38,7 +38,7 @@ blueprint = connection('arg1', 'arg2', kwarg='value') You can link multiple blueprints together with `autoconnect`: ```python session=blueprint-ex1 -from dimos.core.blueprints import autoconnect +from dimos.core.coordination.blueprints import autoconnect class Config(ModuleConfig): arg1: int = 42 @@ -105,7 +105,7 @@ Imagine you have this code: ```python session=blueprint-ex1 from functools import partial -from dimos.core.blueprints import Blueprint, autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out, In @@ -168,7 +168,7 @@ Note: `expanded_blueprint` does not get the transport overrides because it's cre Sometimes you need to rename a connection to match what other modules expect. You can use `remappings` to rename module connections: ```python session=blueprint-ex2 -from dimos.core.blueprints import autoconnect +from dimos.core.coordination.blueprints import autoconnect from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.stream import Out, In @@ -322,6 +322,34 @@ class SomeSkill(Module): return "result" ``` +## Recording + +Blueprints can declare which modules are sensor sources via `default_record_modules`. When `--record-path` is used, only these modules' Out streams are recorded: + +```python session=blueprint-ex4 +from dimos.robot.unitree.go2.connection import GO2Connection + +unitree_go2 = ( + autoconnect(with_vis, GO2Connection.blueprint(), WebsocketVisModule.blueprint()) + .global_config(n_workers=4, robot_model="unitree_go2") + .default_record_modules(GO2Connection) +) +``` + +If `default_record_modules` is not set, all modules with Out streams are recorded. + +Recording captures decoded message objects directly via `Out.subscribe()` — no transport dependency. Codecs are auto-selected (JPEG for images, LCM for standard messages). Sensor timestamps are preserved from `msg.ts`. + +```bash +# Record sensor outputs +dimos --record-path recording.db run unitree-go2 + +# Replay (disables modules whose Outs are in the recording) +dimos --replay-file recording.db --viewer rerun-web run unitree-go2 +``` + +During replay, modules whose Out stream names overlap with the recording are disabled. The recorded data is published to LCM at realtime speed, and `rerun-bridge` handles visualization. + ## Building All you have to do to build a blueprint is call: diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 38366f5daf..6b5f561f7c 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -17,6 +17,8 @@ dimos [GLOBAL OPTIONS] COMMAND [ARGS] | `--simulation` / `--no-simulation` | bool | `False` | Enable MuJoCo simulation | | `--replay` / `--no-replay` | bool | `False` | Use recorded replay data | | `--replay-dir` | TEXT | `go2_sf_office` | Replay dataset directory name | +| `--replay-file` | TEXT | `None` | Path to a `.db` recording file for replay | +| `--record-path` | TEXT | `None` | Record module outputs to a `.db` file | | `--new-memory` / `--no-new-memory` | bool | `False` | Clear persistent memory on start | | `--viewer` | `rerun\|rerun-web\|rerun-connect\|foxglove\|none` | `rerun` | Visualization backend | | `--n-workers` | INT | `2` | Number of forkserver workers | @@ -75,9 +77,15 @@ dimos run unitree-go2 # Background (returns immediately) dimos run unitree-go2-agentic --daemon -# Replay with Rerun viewer +# Replay with Rerun viewer (legacy pickle datasets) dimos --replay --viewer rerun run unitree-go2 +# Record all sensor outputs to a database +dimos --record-path recording.db run unitree-go2 + +# Replay a recording (disables recorded modules, replays via LCM) +dimos --replay-file recording.db --viewer rerun-web run unitree-go2 + # Real robot dimos run unitree-go2-agentic --robot-ip 192.168.123.161 @@ -310,6 +318,15 @@ Or run standalone: dtop ``` +### `recorder` + +Record and replay tool — a terminal UI for recording LCM streams, playing back recordings, and trimming. +It allows interactive selection of topics to record. + +```bash +dimos recorder [filename] +``` + ### `rerun-bridge` Launch the Rerun visualization bridge as a standalone process (outside of a blueprint). From 4e6ea2c61f3cbb93debe61325f50522f5318adad Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 15:49:31 +0100 Subject: [PATCH 14/33] Typing fixes --- .../examples/twist_base_keyboard_teleop.py | 8 ++-- dimos/core/coordination/module_coordinator.py | 11 ++--- dimos/memory2/observationstore/memory.py | 4 +- dimos/memory2/stream.py | 2 + dimos/protocol/pubsub/spec.py | 7 +-- dimos/utils/cli/recorder/run_recorder.py | 44 +++++++------------ dimos/visualization/rerun/bridge.py | 25 ++++++----- 7 files changed, 45 insertions(+), 56 deletions(-) diff --git a/dimos/control/examples/twist_base_keyboard_teleop.py b/dimos/control/examples/twist_base_keyboard_teleop.py index 43b2270410..8c243c8124 100644 --- a/dimos/control/examples/twist_base_keyboard_teleop.py +++ b/dimos/control/examples/twist_base_keyboard_teleop.py @@ -33,12 +33,14 @@ from __future__ import annotations +import asyncio + from dimos.control.blueprints.mobile import coordinator_mock_twist_base from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop -def main() -> None: +async def main() -> None: """Run mock twist base + keyboard teleop.""" coord = ModuleCoordinator.build(coordinator_mock_twist_base) teleop = ModuleCoordinator.build(KeyboardTeleop.blueprint()) @@ -52,9 +54,9 @@ def main() -> None: teleop.start() # Block until Ctrl+C — loop() handles KeyboardInterrupt and calls stop() - coord.loop() + await coord.loop() teleop.stop() if __name__ == "__main__": - main() + asyncio.run(main()) diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index 5a36c2442e..80f7064803 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -242,11 +242,6 @@ def build( # Strip replay_file from all override sources so the nested build() # inside replay() does not re-enter this branch. global_config.replay_file = None - clean_cli = ( - {k: v for k, v in cli_config_overrides.items() if k != "replay_file"} - if cli_config_overrides - else None - ) clean_bp = dataclasses.replace( blueprint, global_config_overrides=MappingProxyType( @@ -257,7 +252,7 @@ def build( } ), ) - return cls.replay(clean_bp, replay_file, cli_config_overrides=clean_cli) + return cls.replay(clean_bp, replay_file, blueprint_args=dict(blueprint_args)) _run_configurators(blueprint) _check_requirements(blueprint) @@ -295,7 +290,7 @@ def replay( recording_path: str, *, speed: float = 1.0, - cli_config_overrides: Mapping[str, Any] | None = None, + blueprint_args: MutableMapping[str, Any] | None = None, ) -> ModuleCoordinator: """Build with a recording replacing some module outputs.""" recording = RecordReplay(recording_path) @@ -330,7 +325,7 @@ def replay( ) patched = blueprint.disabled_modules(*modules_to_disable) - coordinator = cls.build(patched, cli_config_overrides) + coordinator = cls.build(patched, blueprint_args) recording.play(speed=speed) diff --git a/dimos/memory2/observationstore/memory.py b/dimos/memory2/observationstore/memory.py index 76a20cf5ba..820a20ab83 100644 --- a/dimos/memory2/observationstore/memory.py +++ b/dimos/memory2/observationstore/memory.py @@ -92,5 +92,7 @@ def delete_range(self, t1: float, t2: float) -> list[int]: with self._lock: to_delete = [obs for obs in self._observations if t1 <= obs.ts <= t2] ids = [obs.id for obs in to_delete] - self._observations = [obs for obs in self._observations if not (t1 <= obs.ts <= t2)] + self._observations = deque( + obs for obs in self._observations if not (t1 <= obs.ts <= t2) + ) return ids diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index a305bd54e4..c6776045a3 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -356,6 +356,8 @@ def delete_range(self, t1: float, t2: float) -> int: """Delete all observations with timestamps in [t1, t2]. Returns count deleted.""" if isinstance(self._source, Stream): raise TypeError("Cannot delete from a transform stream.") + if self._source is None: + raise TypeError("No source available.") return self._source.delete_range(t1, t2) def publish(self, out: Any) -> DisposableBase: diff --git a/dimos/protocol/pubsub/spec.py b/dimos/protocol/pubsub/spec.py index 63709b39e0..0e292cfcb7 100644 --- a/dimos/protocol/pubsub/spec.py +++ b/dimos/protocol/pubsub/spec.py @@ -17,13 +17,10 @@ from collections.abc import AsyncIterator, Callable from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable - -if TYPE_CHECKING: - from dimos.protocol.pubsub.impl.lcmpubsub import Topic +from typing import Any, Generic, Protocol, TypeVar, runtime_checkable MsgT = TypeVar("MsgT") -TopicT = TypeVar("TopicT", bound="Topic") +TopicT = TypeVar("TopicT") MsgT_co = TypeVar("MsgT_co", covariant=True) TopicT_co = TypeVar("TopicT_co", covariant=True) diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index f743b2fc62..0ed3fd2a82 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -27,7 +27,7 @@ from __future__ import annotations from collections import deque -from pathlib import Path +import sys import time from typing import Any @@ -234,11 +234,7 @@ def on_mount(self) -> None: self._spy = GraphLCMSpy(graph_log_window=0.5) self._spy.start() - if self._db_path: - self._recorder = RecordReplay(self._db_path) - else: - self._recorder = RecordReplay() - + self._recorder = RecordReplay(self._db_path) self.title = f"recorder — {self._recorder.path}" self.set_interval(0.5, self._refresh) @@ -390,7 +386,7 @@ async def _seek_relative(self, delta: float) -> None: def _refresh(self) -> None: if self._table is None: return - rec = self._recorder + assert self._recorder is not None spy = self._spy # ---- Build unified row list: live topics + recorded-only streams ---- @@ -405,7 +401,7 @@ def _refresh(self) -> None: rows[sname] = (channel, spy_topic) # Add streams that exist in the recording but are not live - recorded_streams = set(rec.store.list_streams()) if rec else set() + recorded_streams = set(self._recorder.store.list_streams()) for sname in recorded_streams: if sname not in rows: rows[sname] = (sname, None) @@ -427,7 +423,7 @@ def _refresh(self) -> None: sel = Text(" [ ] ", style=theme.DIM) # Topic name — green when actively recording, bright when selected - if rec and rec.is_recording and sname in (rec.store.list_streams()): + if self._recorder.is_recording and sname in (self._recorder.store.list_streams()): topic_style = f"bold {theme.BRIGHT_GREEN}" elif is_sel: topic_style = theme.BRIGHT_WHITE @@ -449,11 +445,8 @@ def _refresh(self) -> None: # Recorded count if sname in recorded_streams: - try: - count = rec.stream(sname).count() - rec_text = Text(str(count), style=theme.YELLOW) - except Exception: - rec_text = Text("—", style=theme.DIM) + count = self._recorder.store.stream(sname).count() + rec_text = Text(str(count), style=theme.YELLOW) else: rec_text = Text("", style=theme.DIM) @@ -484,8 +477,8 @@ def _refresh(self) -> None: self._table.move_cursor(row=row) # ---- Timeline ---- - duration = rec.duration if rec else 0.0 - position = rec.position if rec else 0.0 + duration = self._recorder.duration + position = self._recorder.position timeline = Text() timeline.append(" ") @@ -513,11 +506,11 @@ def _refresh(self) -> None: # ---- Status bar ---- status = Text() - if rec and rec.is_recording: + if self._recorder.is_recording: status.append(" ● REC ", style=f"bold on {theme.RED}") - elif rec and rec.is_paused: + elif self._recorder.is_paused: status.append(" ❚❚ PAUSED ", style=theme.YELLOW) - elif rec and rec.is_playing: + elif self._recorder.is_playing: status.append(" ▶ PLAYING ", style=theme.BRIGHT_GREEN) else: status.append(" ■ STOPPED ", style=theme.DIM) @@ -531,7 +524,7 @@ def _refresh(self) -> None: status.append(f" {len(recorded_streams)} recorded", style=theme.YELLOW) # Contextual hint - if not (rec and (rec.is_recording or rec.is_playing)): + if not (self._recorder.is_recording or self._recorder.is_playing): if n_live > 0 and n_sel == 0: status.append(" SPACE select, A all, R rec", style=theme.DIM) elif n_sel > 0: @@ -540,23 +533,20 @@ def _refresh(self) -> None: self.query_one("#status-left", Static).update(status) rhs = Text() - if rec: - rhs.append(f"{rec.path} ", style=theme.DIM) + rhs.append(f"{self._recorder.path} ", style=theme.DIM) self.query_one("#status-right", Static).update(rhs) def main() -> None: - import sys - - db_path: Path | None = None + db_path: str | None = None autoplay = False args = sys.argv[1:] if args and args[0] == "play" and len(args) > 1: - db_path = Path(args[1]) + db_path = args[1] autoplay = True elif args and not args[0].startswith("-"): - db_path = Path(args[0]) + db_path = args[0] RecorderApp(db_path=db_path, autoplay=autoplay).run() diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index dd2dbfc361..c91f7e1f60 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -16,17 +16,17 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import field from functools import lru_cache import subprocess +import sys import time from typing import ( Any, Literal, Protocol, TypeAlias, - TypeGuard, cast, runtime_checkable, ) @@ -46,6 +46,11 @@ from dimos.protocol.pubsub.spec import SubscribeAllCapable from dimos.utils.logging_config import setup_logger +if sys.version_info >= (3, 13): + from typing import TypeIs +else: + from typing import TypeGuard as TypeIs + RERUN_GRPC_PORT = 9876 RERUN_WEB_PORT = 9090 @@ -95,20 +100,16 @@ BlueprintFactory: TypeAlias = Callable[[], "Blueprint"] # to_rerun() can return a single archetype or a list of (entity_path, archetype) tuples -RerunMulti: TypeAlias = "list[tuple[str, Archetype]]" +RerunMulti: TypeAlias = Sequence[tuple[str, Archetype]] RerunData: TypeAlias = "Archetype | RerunMulti" -def is_rerun_multi(data: Any) -> TypeGuard[RerunMulti]: +def is_rerun_multi(data: object) -> TypeIs[RerunMulti]: """Check if data is a list of (entity_path, archetype) tuples.""" - return ( - isinstance(data, list) - and bool(data) - and isinstance(data[0], tuple) - and len(data[0]) == 2 - and isinstance(data[0][0], str) - and isinstance(data[0][1], Archetype) - ) + match data: + case [(str(), Archetype()), *_]: + return True + return False @runtime_checkable From ad6debf9998c204962c93a17817fc92fff801833 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 16:42:58 +0100 Subject: [PATCH 15/33] Typing fixes --- dimos/record/record_replay.py | 2 +- dimos/visualization/rerun/bridge.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index 025f84aafd..b0b345bba6 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -310,7 +310,7 @@ async def _playback_loop(self) -> None: for path, archetype in rerun_data: rec.log(path, archetype) else: - rec.log(entity_path, rerun_data) + rec.log(entity_path, rerun_data) # type: ignore[arg-type] try: next_obs = next(it) diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index c91f7e1f60..437c03944f 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -16,7 +16,7 @@ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Callable from dataclasses import field from functools import lru_cache import subprocess @@ -100,7 +100,7 @@ BlueprintFactory: TypeAlias = Callable[[], "Blueprint"] # to_rerun() can return a single archetype or a list of (entity_path, archetype) tuples -RerunMulti: TypeAlias = Sequence[tuple[str, Archetype]] +RerunMulti: TypeAlias = list[tuple[str, Archetype]] RerunData: TypeAlias = "Archetype | RerunMulti" From b11e9383ba63c2de99481ab03bb7ab94afd05fa7 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 16:45:40 +0100 Subject: [PATCH 16/33] Revert --- dimos/types/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 dimos/types/__init__.py diff --git a/dimos/types/__init__.py b/dimos/types/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 From bc563b73e035bd820f2c0192c09d7a5d968c8d46 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 17:09:25 +0100 Subject: [PATCH 17/33] Remove init files --- dimos/record/__init__.py | 17 ----------------- dimos/utils/cli/recorder/__init__.py | 0 2 files changed, 17 deletions(-) delete mode 100644 dimos/record/__init__.py delete mode 100644 dimos/utils/cli/recorder/__init__.py diff --git a/dimos/record/__init__.py b/dimos/record/__init__.py deleted file mode 100644 index 1bb0bfbcc9..0000000000 --- a/dimos/record/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dimos.record.record_replay import RecordReplay - -__all__ = ("RecordReplay",) diff --git a/dimos/utils/cli/recorder/__init__.py b/dimos/utils/cli/recorder/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 From 2d6f1590736d77a2b2ccb0983855fa8f0e346efe Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 17:29:46 +0100 Subject: [PATCH 18/33] Fix tests --- dimos/record/test_record_replay.py | 45 ++++++++++++++++-------- dimos/utils/cli/recorder/run_recorder.py | 20 +++-------- 2 files changed, 35 insertions(+), 30 deletions(-) diff --git a/dimos/record/test_record_replay.py b/dimos/record/test_record_replay.py index f75ef121a2..bd75ee9138 100644 --- a/dimos/record/test_record_replay.py +++ b/dimos/record/test_record_replay.py @@ -45,7 +45,7 @@ class FakePubSub: """Minimal PubSub that supports subscribe_all for testing.""" def __init__(self) -> None: - self._subscribers: list[Callable[[Any, Any], Any]] = [] + self._subscribers: list[Callable[[Any, Any], None]] = [] self._lock = threading.Lock() self._started = False @@ -59,14 +59,20 @@ def publish(self, topic: Any, message: Any) -> None: # Not needed for recording tests pass - def subscribe_all(self, callback: Callable[[Any, Any], Any]) -> Callable[[], None]: + def subscribe(self, topic: object, callback: Callable[[Any, Any], None]) -> Callable[[], None]: + key = str(topic) + + def filtered(msg: Any, t: Any) -> None: + if str(t) == key: + callback(msg, t) + with self._lock: - self._subscribers.append(callback) + self._subscribers.append(filtered) def unsub() -> None: with self._lock: with suppress(ValueError): - self._subscribers.remove(callback) + self._subscribers.remove(filtered) return unsub @@ -89,6 +95,17 @@ def __eq__(self, other: object) -> bool: return isinstance(other, SimpleMsg) and self.value == other.value +ALL_TOPICS = [ + FakeTopic("/sensor/lidar"), + FakeTopic("/sensor/odom"), + FakeTopic("/sensor"), + FakeTopic("/data"), + FakeTopic("/a"), + FakeTopic("/from_ps1"), + FakeTopic("/from_ps2"), +] + + @pytest.fixture def tmp_db(tmp_path: Path) -> str: return str(tmp_path / "test_recording.db") @@ -98,7 +115,7 @@ class TestRecordReplay: async def test_record_and_list_streams(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) assert rec.is_recording pubsub.emit("/sensor/lidar", SimpleMsg(1.0)) @@ -116,7 +133,7 @@ async def test_record_and_list_streams(self, tmp_db: str) -> None: async def test_record_and_query(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(10): pubsub.emit("/data", SimpleMsg(float(i))) @@ -133,7 +150,7 @@ async def test_record_and_query(self, tmp_db: str) -> None: async def test_duration(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) pubsub.emit("/a", SimpleMsg(0.0)) await asyncio.sleep(0.1) pubsub.emit("/a", SimpleMsg(1.0)) @@ -144,7 +161,7 @@ async def test_duration(self, tmp_db: str) -> None: async def test_stream_info(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(5): pubsub.emit("/sensor", SimpleMsg(float(i))) await asyncio.sleep(0.01) @@ -158,7 +175,7 @@ async def test_stream_info(self, tmp_db: str) -> None: async def test_delete_range(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(20): pubsub.emit("/data", SimpleMsg(float(i))) await asyncio.sleep(0.01) @@ -178,7 +195,7 @@ async def test_delete_range(self, tmp_db: str) -> None: async def test_trim(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(30): pubsub.emit("/data", SimpleMsg(float(i))) await asyncio.sleep(0.01) @@ -202,7 +219,7 @@ async def test_playback_runs(self, tmp_db: str) -> None: """Test that playback task starts and finishes.""" pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(5): pubsub.emit("/data", SimpleMsg(float(i))) await asyncio.sleep(0.01) @@ -217,7 +234,7 @@ async def test_playback_runs(self, tmp_db: str) -> None: async def test_stop_playback(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(100): pubsub.emit("/data", SimpleMsg(float(i))) await asyncio.sleep(0.005) @@ -232,7 +249,7 @@ async def test_stop_playback(self, tmp_db: str) -> None: async def test_seek(self, tmp_db: str) -> None: pubsub = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([pubsub]) + rec.start_recording([pubsub], topics=ALL_TOPICS) for i in range(10): pubsub.emit("/data", SimpleMsg(float(i))) await asyncio.sleep(0.01) @@ -245,7 +262,7 @@ async def test_multiple_pubsubs(self, tmp_db: str) -> None: ps1 = FakePubSub() ps2 = FakePubSub() async with RecordReplay(tmp_db) as rec: - rec.start_recording([ps1, ps2]) + rec.start_recording([ps1, ps2], topics=ALL_TOPICS) ps1.emit("/from_ps1", SimpleMsg(1.0)) ps2.emit("/from_ps2", SimpleMsg(2.0)) rec.stop_recording() diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index 0ed3fd2a82..8616bdfd7d 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -249,10 +249,6 @@ async def on_unmount(self) -> None: if self._lcm and hasattr(self._lcm, "stop"): self._lcm.stop() - # ------------------------------------------------------------------ - # Actions - # ------------------------------------------------------------------ - def action_toggle_select(self) -> None: """Toggle selection on the row under the cursor.""" if self._table is None or self._table.row_count == 0: @@ -343,10 +339,6 @@ async def action_do_delete(self) -> None: self._recorder.delete_range(lo, hi) self._trim_in = self._trim_out = None - # ------------------------------------------------------------------ - # Helpers - # ------------------------------------------------------------------ - def _all_topics(self) -> list[Any]: """All discovered topics from the spy.""" from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic @@ -379,17 +371,13 @@ async def _seek_relative(self, delta: float) -> None: if self._recorder: await self._recorder.seek(self._recorder.position + delta) - # ------------------------------------------------------------------ - # Refresh - # ------------------------------------------------------------------ - def _refresh(self) -> None: if self._table is None: return assert self._recorder is not None spy = self._spy - # ---- Build unified row list: live topics + recorded-only streams ---- + # Build unified row list: live topics + recorded-only streams # Each row: (stream_name, channel, spy_topic_or_None) rows: dict[str, tuple[str, Any]] = {} # stream_name -> (channel, spy_topic) @@ -406,7 +394,7 @@ def _refresh(self) -> None: if sname not in rows: rows[sname] = (sname, None) - # ---- Render table ---- + # Render table # Remember cursor position so we can restore it cursor_row = self._table.cursor_coordinate.row if self._table.row_count > 0 else 0 self._table.clear(columns=False) @@ -476,7 +464,7 @@ def _refresh(self) -> None: row = min(cursor_row, self._table.row_count - 1) self._table.move_cursor(row=row) - # ---- Timeline ---- + # Timeline duration = self._recorder.duration position = self._recorder.position @@ -504,7 +492,7 @@ def _refresh(self) -> None: self.query_one("#timeline", Static).update(timeline) - # ---- Status bar ---- + # Status bar status = Text() if self._recorder.is_recording: status.append(" ● REC ", style=f"bold on {theme.RED}") From 5b4bdf4be1d83e14cf17ddd596f6cd5946d8f501 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 17:43:24 +0100 Subject: [PATCH 19/33] Fix --- dimos/core/coordination/module_coordinator.py | 2 +- dimos/utils/cli/recorder/run_recorder.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index 80f7064803..4033c09e3d 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -32,7 +32,7 @@ from dimos.core.module import ModuleBase, ModuleSpec from dimos.core.resource import Resource from dimos.core.transport import LCMTransport, PubSubTransport, pLCMTransport -from dimos.record import RecordReplay +from dimos.record.record_replay import RecordReplay from dimos.spec.utils import spec_annotation_compliance, spec_structural_compliance from dimos.utils.generic import short_id from dimos.utils.logging_config import setup_logger diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index 8616bdfd7d..03f878b7eb 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -225,7 +225,6 @@ def compose(self) -> ComposeResult: def on_mount(self) -> None: from dimos.protocol.pubsub.impl.lcmpubsub import LCM - from dimos.record import RecordReplay from dimos.utils.cli.lcmspy.lcmspy import GraphLCMSpy self._lcm = LCM() From 2dfc23ee01a1606ca833fed0520178ac18bd048b Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 17:55:12 +0100 Subject: [PATCH 20/33] Fix --- dimos/record/test_record_replay.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/record/test_record_replay.py b/dimos/record/test_record_replay.py index bd75ee9138..1223eba7ee 100644 --- a/dimos/record/test_record_replay.py +++ b/dimos/record/test_record_replay.py @@ -23,7 +23,7 @@ import pytest -from dimos.record import RecordReplay +from dimos.record.record_replay import RecordReplay class FakeTopic: From 787bd6625aadeccde8e8adbc8a224ae4d06474a6 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 19:05:09 +0100 Subject: [PATCH 21/33] Fix --- dimos/core/test_core.py | 2 +- dimos/record/test_record_replay.py | 2 +- dimos/robot/test_all_blueprints_generation.py | 9 ++++++++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/dimos/core/test_core.py b/dimos/core/test_core.py index bf62e526ee..e69cde2fef 100644 --- a/dimos/core/test_core.py +++ b/dimos/core/test_core.py @@ -77,7 +77,7 @@ def test_classmethods() -> None: # Check that we have the expected RPC methods assert "navigate_to" in class_rpcs, "navigate_to should be in rpcs" assert "start" in class_rpcs, "start should be in rpcs" - assert len(class_rpcs) == 7 + assert len(class_rpcs) == 9 # Check that the values are callable assert callable(class_rpcs["navigate_to"]), "navigate_to should be callable" diff --git a/dimos/record/test_record_replay.py b/dimos/record/test_record_replay.py index 1223eba7ee..adac904c1d 100644 --- a/dimos/record/test_record_replay.py +++ b/dimos/record/test_record_replay.py @@ -227,7 +227,7 @@ async def test_playback_runs(self, tmp_db: str) -> None: rec.play(speed=100.0) # very fast assert rec.is_playing - async with asyncio.timeout(0.1): + async with asyncio.timeout(5.0): await rec._play_task assert not rec.is_playing diff --git a/dimos/robot/test_all_blueprints_generation.py b/dimos/robot/test_all_blueprints_generation.py index d40ad2aed5..45aa7855c1 100644 --- a/dimos/robot/test_all_blueprints_generation.py +++ b/dimos/robot/test_all_blueprints_generation.py @@ -32,7 +32,14 @@ "dimos/core/blueprints.py", "dimos/core/test_blueprints.py", } -BLUEPRINT_METHODS = {"transports", "global_config", "remappings", "requirements", "configurators"} +BLUEPRINT_METHODS = { + "transports", + "global_config", + "remappings", + "requirements", + "configurators", + "default_record_modules", +} _EXCLUDED_MODULE_NAMES = {"Module", "ModuleBase", "StreamModule"} From d9f6ff392f940c621406613d3ae0be36e88253dc Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 19:47:43 +0100 Subject: [PATCH 22/33] Adjust imports --- dimos/core/module.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 2bacca16d0..c021c6e6df 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -41,7 +41,6 @@ from dimos.core.rpc_client import RpcCall from dimos.core.stream import In, Out, RemoteOut, Transport from dimos.memory2.store.sqlite import SqliteStore -from dimos.memory2.stream import Stream from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable @@ -54,6 +53,7 @@ from dimos.core.coordination.blueprints import Blueprint from dimos.core.introspection.module.info import ModuleInfo from dimos.core.rpc_client import RPCClient + from dimos.memory2.stream import Stream if sys.version_info >= (3, 13): from typing import TypeVar @@ -385,6 +385,8 @@ def blueprint(self) -> _BlueprintPartial: @rpc def start_recording(self, db_path: str) -> None: + from dimos.memory2.store.sqlite import SqliteStore + self._rec_store = SqliteStore(path=db_path) for name, out in self.outputs.items(): stream = self._rec_store.stream(name, out.type) @@ -393,7 +395,7 @@ def start_recording(self, db_path: str) -> None: reg["channel"] = f"/{name}#{out.type.msg_name}" self._rec_store._registry.put(name, reg) - def cb(msg: Any, _stream: Stream[object] = stream) -> None: + def cb(msg: Any, _stream: "Stream[object]" = stream) -> None: ts = msg.ts if isinstance(msg, Timestamped) else time.time() _stream.append(msg, ts=ts) From 480f853028e7877b5bc7025a14f85019a66d945f Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Mon, 13 Apr 2026 22:11:30 +0100 Subject: [PATCH 23/33] Fix args --- dimos/robot/cli/dimos.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index 9fc30cc8a0..f12593eb68 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -195,14 +195,12 @@ def load_config_args(config: type[BaseModel], args: Iterable[str], path: Path) - async def _run( ctx: typer.Context, - robot_types: list[str] = typer.Argument(..., help="Blueprints or modules to run"), - daemon: bool = typer.Option(False, "--daemon", "-d", help="Run in background"), - disable: list[str] = typer.Option([], "--disable", help="Module names to disable"), - blueprint_args: list[str] = typer.Option((), "--option", "-o"), - config_path: Path = typer.Option( - CONFIG_DIR / "dimos", "--config", "-c", help="Path to config file" - ), - show_help: bool = typer.Option(False, "--help"), + robot_types: list[str], + daemon: bool, + disable: list[str], + blueprint_args: list[str], + config_path: Path, + show_help: bool, ) -> None: """Start a robot blueprint""" logger.info("Starting DimOS") @@ -325,10 +323,15 @@ def run( robot_types: list[str] = typer.Argument(..., help="Blueprints or modules to run"), daemon: bool = typer.Option(False, "--daemon", "-d", help="Run in background"), disable: list[str] = typer.Option([], "--disable", help="Module names to disable"), + blueprint_args: list[str] = typer.Option((), "--option", "-o"), + config_path: Path = typer.Option( + CONFIG_DIR / "dimos", "--config", "-c", help="Path to config file" + ), + show_help: bool = typer.Option(False, "--help"), ) -> None: """Start a robot blueprint""" logger.info("Starting DimOS") - asyncio.run(_run(ctx, robot_types, daemon, disable)) + asyncio.run(_run(ctx, robot_types, daemon, disable, blueprint_args, config_path, show_help)) @main.command() From 28de1d252cdb55e7fa1c0571d0d00e5ff23529a4 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Tue, 14 Apr 2026 00:06:32 +0100 Subject: [PATCH 24/33] Fixes --- dimos/robot/cli/dimos.py | 5 +++++ dimos/utils/cli/recorder/run_recorder.py | 4 ++-- dimos/visualization/rerun/bridge.py | 7 ++++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index f12593eb68..842883eb5f 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -220,6 +220,11 @@ async def _run( cli_config_overrides: dict[str, Any] = ctx.obj + # Apply overrides early so blueprint module-level code (e.g. viewer + # selection) sees the correct global_config values at import time. + if cli_config_overrides: + global_config.update(**cli_config_overrides) + # Clean stale registry entries stale = cleanup_stale() if stale: diff --git a/dimos/utils/cli/recorder/run_recorder.py b/dimos/utils/cli/recorder/run_recorder.py index 03f878b7eb..0dbc4db5a0 100644 --- a/dimos/utils/cli/recorder/run_recorder.py +++ b/dimos/utils/cli/recorder/run_recorder.py @@ -344,7 +344,7 @@ def _all_topics(self) -> list[Any]: if not self._spy: return [] - return [LCMTopic.from_channel_str(ch) for ch in self._spy.topic] + return [LCMTopic.from_channel_str(ch) for ch in list(self._spy.topic)] def _selected_topics(self) -> list[Any]: """Map selected stream names back to LCM Topics via the spy.""" @@ -355,7 +355,7 @@ def _selected_topics(self) -> list[Any]: return [] return [ LCMTopic.from_channel_str(ch) - for ch in self._spy.topic + for ch in list(self._spy.topic) if topic_to_stream_name(ch) in self._selected ] diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 437c03944f..07ca34bb5a 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -307,6 +307,7 @@ def start(self) -> None: rr.init("dimos") if self.config.viewer_mode == "native": + dimos_viewer = False try: import rerun_bindings @@ -315,6 +316,7 @@ def start(self) -> None: executable_name="dimos-viewer", memory_limit=self.config.memory_limit, ) + dimos_viewer = True except ImportError: pass # dimos-viewer not installed except Exception: @@ -322,7 +324,10 @@ 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 dimos_viewer: + rr.connect_grpc(f"rerun+http://127.0.0.1:{RERUN_GRPC_PORT}/proxy") + else: + rr.spawn(connect=True, memory_limit=self.config.memory_limit) elif self.config.viewer_mode == "web": server_uri = rr.serve_grpc(grpc_port=RERUN_GRPC_PORT) rr.serve_web_viewer(connect_to=server_uri, open_browser=False) From 34a20baabc64afd32baf3c72b5f6af0d20395ebc Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Tue, 14 Apr 2026 01:09:49 +0100 Subject: [PATCH 25/33] Fix test --- dimos/record/test_record_replay.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/dimos/record/test_record_replay.py b/dimos/record/test_record_replay.py index adac904c1d..08124f4870 100644 --- a/dimos/record/test_record_replay.py +++ b/dimos/record/test_record_replay.py @@ -20,6 +20,7 @@ from pathlib import Path import threading from typing import Any +from unittest.mock import patch import pytest @@ -225,11 +226,12 @@ async def test_playback_runs(self, tmp_db: str) -> None: await asyncio.sleep(0.01) rec.stop_recording() - rec.play(speed=100.0) # very fast - assert rec.is_playing - async with asyncio.timeout(5.0): - await rec._play_task - assert not rec.is_playing + with patch("dimos.record.record_replay.rr", autospec=True, spec_set=True): + rec.play(speed=100.0) # very fast + assert rec.is_playing + async with asyncio.timeout(5.0): + await rec._play_task + assert not rec.is_playing async def test_stop_playback(self, tmp_db: str) -> None: pubsub = FakePubSub() @@ -240,10 +242,11 @@ async def test_stop_playback(self, tmp_db: str) -> None: await asyncio.sleep(0.005) rec.stop_recording() - rec.play(speed=0.1) # slow - await asyncio.sleep(0.1) - assert rec.is_playing - await rec.stop_playback() + with patch("dimos.record.record_replay.rr", autospec=True, spec_set=True): + rec.play(speed=0.1) # slow + await asyncio.sleep(0.1) + assert rec.is_playing + await rec.stop_playback() assert not rec.is_playing async def test_seek(self, tmp_db: str) -> None: From c3588c02798196e009afcba22f97f9082754a7a8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 21 Apr 2026 09:48:06 +0000 Subject: [PATCH 26/33] docs: add record/replay usage and test steps Agent-Logs-Url: https://github.com/dimensionalOS/dimos/sessions/44dd2709-d3ba-4138-b140-56e76fd15eb9 Co-authored-by: leshy <681516+leshy@users.noreply.github.com> --- docs/usage/cli.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/docs/usage/cli.md b/docs/usage/cli.md index ff0a2dff67..9dae239274 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -99,6 +99,24 @@ dimos run unitree-go2 keyboard-teleop dimos run unitree-go2-agentic --disable OsmSkill WebInput ``` +#### Record/replay end-to-end test + +```bash +# 1) Start a stack and record all configured output streams +dimos --record-path /tmp/recording.db run unitree-go2 + +# 2) Stop the run, then replay the same file +dimos --replay-file /tmp/recording.db --viewer rerun-web run unitree-go2 +``` + +For interactive topic selection, use `dimos recorder` instead of `--record-path`: + +```bash +dimos recorder /tmp/recording.db +``` + +Then select streams with `space` and press `r` to start/stop recording. + When `--daemon` is used, the process: 1. Builds and starts all modules (foreground — you see errors) 2. Runs a health check (polls worker PIDs) @@ -330,6 +348,15 @@ It allows interactive selection of topics to record. dimos recorder [filename] ``` +Common controls: +- `space`: select/deselect current stream +- `r`: start/stop recording selected streams +- `p`: play/pause recording +- `s`: stop playback/recording +- `[` / `]`: mark trim in/out +- `t`: trim to marked range +- `d`: delete marked range + ### `rerun-bridge` Launch the Rerun visualization bridge as a standalone process (outside of a blueprint). From d0f1b63584b500fd578c166f9aac816f8d1d78c8 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 14:53:00 +0800 Subject: [PATCH 27/33] adressing Paul pr comments --- dimos/core/coordination/module_coordinator.py | 3 ++- dimos/core/module.py | 19 +++++++-------- dimos/memory2/observationstore/base.py | 3 ++- dimos/memory2/observationstore/sqlite.py | 4 ++-- dimos/memory2/store/sqlite.py | 14 +++++++++++ dimos/record/record_replay.py | 23 +++++++++---------- dimos/robot/cli/dimos.py | 2 -- docs/usage/blueprints.md | 2 -- uv.lock | 2 -- 9 files changed, 41 insertions(+), 31 deletions(-) diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index 2543badf70..a6b9d6f744 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -22,6 +22,7 @@ from pathlib import Path import shutil import sys +import threading from types import MappingProxyType from typing import TYPE_CHECKING, Any, cast @@ -302,7 +303,7 @@ def build( if bp.module in record_modules: instance = coordinator.get_instance(bp.module) if instance is not None: - instance.start_recording(global_config.record_path) + instance.start_recording_outputs(global_config.record_path) _log_blueprint_graph(blueprint, coordinator) diff --git a/dimos/core/module.py b/dimos/core/module.py index c021c6e6df..b08eef2207 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -161,7 +161,7 @@ def start(self) -> None: @rpc def stop(self) -> None: super().stop() - self._stop_recording() + self._stop_recording_outputs() self._close_module() def _close_module(self) -> None: @@ -384,16 +384,17 @@ def blueprint(self) -> _BlueprintPartial: return partial(Blueprint.create, self) # type: ignore[arg-type] @rpc - def start_recording(self, db_path: str) -> None: + def start_recording_outputs(self, db_path: str) -> None: from dimos.memory2.store.sqlite import SqliteStore + if self._rec_store is not None: + raise RuntimeError("Recording already in progress") self._rec_store = SqliteStore(path=db_path) for name, out in self.outputs.items(): stream = self._rec_store.stream(name, out.type) - reg = self._rec_store._registry.get(name) - if reg and "channel" not in reg: - reg["channel"] = f"/{name}#{out.type.msg_name}" - self._rec_store._registry.put(name, reg) + meta = self._rec_store.get_stream_meta(name) + if meta and "channel" not in meta: + self._rec_store.update_stream_meta(name, channel=f"/{name}#{out.type.msg_name}") def cb(msg: Any, _stream: "Stream[object]" = stream) -> None: ts = msg.ts if isinstance(msg, Timestamped) else time.time() @@ -402,10 +403,10 @@ def cb(msg: Any, _stream: "Stream[object]" = stream) -> None: self._rec_unsubs.append(out.subscribe(cb)) @rpc - def stop_recording(self) -> None: - self._stop_recording() + def stop_recording_outputs(self) -> None: + self._stop_recording_outputs() - def _stop_recording(self) -> None: + def _stop_recording_outputs(self) -> None: for unsub in self._rec_unsubs: unsub() self._rec_unsubs = [] diff --git a/dimos/memory2/observationstore/base.py b/dimos/memory2/observationstore/base.py index 392950fe8a..d33fb421c7 100644 --- a/dimos/memory2/observationstore/base.py +++ b/dimos/memory2/observationstore/base.py @@ -69,9 +69,10 @@ def fetch_by_ids(self, ids: list[int]) -> list[Observation[T]]: """Batch fetch by id (for vector search results).""" ... + @abstractmethod def delete_range(self, t1: float, t2: float) -> list[int]: """Delete observations with ts in [t1, t2]. Returns deleted IDs.""" - raise NotImplementedError + ... def serialize(self) -> dict[str, Any]: return {"class": qual(type(self)), "config": self.config.model_dump()} diff --git a/dimos/memory2/observationstore/sqlite.py b/dimos/memory2/observationstore/sqlite.py index 654babbc81..4b91d8d3fd 100644 --- a/dimos/memory2/observationstore/sqlite.py +++ b/dimos/memory2/observationstore/sqlite.py @@ -443,12 +443,12 @@ def delete_range(self, t1: float, t2: float) -> list[int]: """Delete observations with ts in [t1, t2]. Returns deleted IDs.""" with self._lock: rows = self._conn.execute( - f'SELECT id FROM "{self._name}" WHERE ts >= ? AND ts <= ?', (t1, t2) + f'DELETE FROM "{self._name}" WHERE ts >= ? AND ts <= ? RETURNING id', + (t1, t2), ).fetchall() ids = [r[0] for r in rows] if ids: placeholders = ",".join("?" * len(ids)) - self._conn.execute(f'DELETE FROM "{self._name}" WHERE id IN ({placeholders})', ids) self._conn.execute( f'DELETE FROM "{self._name}_rtree" WHERE id IN ({placeholders})', ids ) diff --git a/dimos/memory2/store/sqlite.py b/dimos/memory2/store/sqlite.py index 55cf6d5777..39ea1410d6 100644 --- a/dimos/memory2/store/sqlite.py +++ b/dimos/memory2/store/sqlite.py @@ -199,6 +199,20 @@ def list_streams(self) -> list[str]: db_names = set(self._registry.list_streams()) return sorted(db_names | set(self._streams.keys())) + def get_stream_meta(self, name: str) -> dict[str, Any] | None: + """Return persisted metadata for a stream, or None if it doesn't exist.""" + return self._registry.get(name) + + def update_stream_meta(self, name: str, **updates: Any) -> None: + """Merge ``updates`` into a stream's persisted metadata. No-op if the + stream has no registry entry yet. + """ + meta = self._registry.get(name) + if meta is None: + return + meta.update(updates) + self._registry.put(name, meta) + def delete_stream(self, name: str) -> None: super().delete_stream(name) self._registry_conn.execute(f'DROP TABLE IF EXISTS "{name}"') diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index b0b345bba6..1ef211765d 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -40,7 +40,6 @@ from collections.abc import Callable, Collection from contextlib import suppress import heapq -import logging import math import re import sys @@ -55,6 +54,7 @@ from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.service.spec import Service from dimos.types.timestamped import Timestamped +from dimos.utils.logging_config import setup_logger from dimos.visualization.rerun.bridge import RerunConvertible, is_rerun_multi if sys.version_info >= (3, 11): @@ -62,7 +62,7 @@ else: from typing import Any as Self -logger = logging.getLogger(__name__) +logger = setup_logger() _SANITIZE_RE = re.compile(r"[^A-Za-z0-9_]") @@ -175,10 +175,9 @@ def _on_message(self, msg: object, topic: Topic) -> None: s = self._store.stream(stream_name, type(msg)) s.append(msg, ts=ts) - reg = self._store._registry.get(stream_name) - if reg and "channel" not in reg: - reg["channel"] = str(topic) - self._store._registry.put(stream_name, reg) + meta = self._store.get_stream_meta(stream_name) + if meta and "channel" not in meta: + self._store.update_stream_meta(stream_name, channel=str(topic)) @property def duration(self) -> float: @@ -221,9 +220,9 @@ def stream_info(self) -> tuple[StreamInfo, ...]: info["end"] = t1 info["duration"] = t1 - t0 # Get payload type from registry - reg = self._store._registry.get(name) - if reg: - info["type"] = reg.get("payload_module", "unknown") + meta = self._store.get_stream_meta(name) + if meta: + info["type"] = meta.get("payload_module", "unknown") result.append(info) return tuple(result) @@ -254,9 +253,9 @@ async def _playback_loop(self) -> None: # Build topic map for decoding raw bytes -> DimosMsg topic_map: dict[str, Topic] = {} for name in self._store.list_streams(): - reg = self._store._registry.get(name) - if reg: - channel = reg.get("channel") + meta = self._store.get_stream_meta(name) + if meta: + channel = meta.get("channel") if channel: topic_map[name] = Topic.from_channel_str(channel) diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index 9883ae2032..0e3b4bbd65 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -204,8 +204,6 @@ async def _run( show_help: bool, ) -> None: """Start a robot blueprint""" - logger.info("Starting DimOS") - from dimos.core.coordination.blueprints import autoconnect from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.coordination.process_lifecycle import ( diff --git a/docs/usage/blueprints.md b/docs/usage/blueprints.md index 90736d06a3..5811b80a14 100644 --- a/docs/usage/blueprints.md +++ b/docs/usage/blueprints.md @@ -375,8 +375,6 @@ unitree_go2 = ( ) ``` -If `default_record_modules` is not set, all modules with Out streams are recorded. - Recording captures decoded message objects directly via `Out.subscribe()` — no transport dependency. Codecs are auto-selected (JPEG for images, LCM for standard messages). Sensor timestamps are preserved from `msg.ts`. ```bash diff --git a/uv.lock b/uv.lock index f78b54a037..ea7c862d09 100644 --- a/uv.lock +++ b/uv.lock @@ -1700,7 +1700,6 @@ dependencies = [ { name = "psutil" }, { name = "pydantic" }, { name = "pydantic-settings" }, - { name = "pytest-asyncio" }, { name = "python-dotenv" }, { name = "pyturbojpeg" }, { name = "reactivex" }, @@ -2080,7 +2079,6 @@ requires-dist = [ { name = "pymavlink", marker = "extra == 'drone'" }, { name = "pyrealsense2", marker = "sys_platform != 'darwin' and extra == 'manipulation'" }, { name = "pytest", marker = "extra == 'dev'", specifier = "==8.3.5" }, - { name = "pytest-asyncio", specifier = ">=0.26.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = "==0.26.0" }, { name = "pytest-env", marker = "extra == 'dev'", specifier = "==1.1.5" }, { name = "pytest-mock", marker = "extra == 'dev'", specifier = "==3.15.0" }, From 2977242ec6f9bc056ae08d18a006c503d4589df5 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 15:35:18 +0800 Subject: [PATCH 28/33] removed store metadata intervention --- dimos/core/module.py | 8 +------- dimos/memory2/store/sqlite.py | 14 -------------- dimos/memory2/stream.py | 10 ++++++++++ dimos/record/record_replay.py | 30 +++++++++++------------------- 4 files changed, 22 insertions(+), 40 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index bd01587eda..d48ff63c33 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -176,7 +176,7 @@ def start(self) -> None: def stop(self) -> None: self._stop_main() super().stop() - self._stop_recording_outputs() + self.stop_recording_outputs() self._close_module() def _close_module(self) -> None: @@ -424,9 +424,6 @@ def start_recording_outputs(self, db_path: str) -> None: self._rec_store = SqliteStore(path=db_path) for name, out in self.outputs.items(): stream = self._rec_store.stream(name, out.type) - meta = self._rec_store.get_stream_meta(name) - if meta and "channel" not in meta: - self._rec_store.update_stream_meta(name, channel=f"/{name}#{out.type.msg_name}") def cb(msg: Any, _stream: "Stream[object]" = stream) -> None: ts = msg.ts if isinstance(msg, Timestamped) else time.time() @@ -436,9 +433,6 @@ def cb(msg: Any, _stream: "Stream[object]" = stream) -> None: @rpc def stop_recording_outputs(self) -> None: - self._stop_recording_outputs() - - def _stop_recording_outputs(self) -> None: for unsub in self._rec_unsubs: unsub() self._rec_unsubs = [] diff --git a/dimos/memory2/store/sqlite.py b/dimos/memory2/store/sqlite.py index 0b85c3ad68..bb2b735c1c 100644 --- a/dimos/memory2/store/sqlite.py +++ b/dimos/memory2/store/sqlite.py @@ -206,20 +206,6 @@ def list_streams(self) -> list[str]: db_names = set(self._registry.list_streams()) return sorted(db_names | set(self._streams.keys())) - def get_stream_meta(self, name: str) -> dict[str, Any] | None: - """Return persisted metadata for a stream, or None if it doesn't exist.""" - return self._registry.get(name) - - def update_stream_meta(self, name: str, **updates: Any) -> None: - """Merge ``updates`` into a stream's persisted metadata. No-op if the - stream has no registry entry yet. - """ - meta = self._registry.get(name) - if meta is None: - return - meta.update(updates) - self._registry.put(name, meta) - def delete_stream(self, name: str) -> None: super().delete_stream(name) self._registry_conn.execute(f'DROP TABLE IF EXISTS "{name}"') diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index a5d9a9b9eb..5e40faf444 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -125,6 +125,16 @@ def is_live(self) -> bool: return self._source.is_live() return False + @property + def data_type(self) -> type: + """Payload type of the root backend, or ``object`` if unbound.""" + current: Any = self + while isinstance(current._source, Stream): + current = current._source + if current._source is None: + return object + return cast("type", current._source.data_type) + def __iter__(self) -> Iterator[O]: if self._source is None: raise TypeError( diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index 1ef211765d..e3e1b021eb 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -107,7 +107,6 @@ def __init__(self, path: str) -> None: self._play_task: asyncio.Task[None] | None = None self._play_speed = 1.0 self._position = 0.0 - self._pubsub = None @property def store(self) -> SqliteStore: @@ -175,10 +174,6 @@ def _on_message(self, msg: object, topic: Topic) -> None: s = self._store.stream(stream_name, type(msg)) s.append(msg, ts=ts) - meta = self._store.get_stream_meta(stream_name) - if meta and "channel" not in meta: - self._store.update_stream_meta(stream_name, channel=str(topic)) - @property def duration(self) -> float: """Total duration of the recording in seconds.""" @@ -219,10 +214,8 @@ def stream_info(self) -> tuple[StreamInfo, ...]: info["start"] = t0 info["end"] = t1 info["duration"] = t1 - t0 - # Get payload type from registry - meta = self._store.get_stream_meta(name) - if meta: - info["type"] = meta.get("payload_module", "unknown") + t = s.data_type + info["type"] = f"{t.__module__}.{t.__qualname__}" if t is not object else "unknown" result.append(info) return tuple(result) @@ -250,14 +243,15 @@ async def _playback_loop(self) -> None: rec = rr.RecordingStream("playback", make_default=False) rec.connect_grpc() - # Build topic map for decoding raw bytes -> DimosMsg + # Reconstruct topic map from each stream's payload type. Channel is + # /#; under the convention that recorded topic + # names are valid identifiers, the stream name is the channel name. topic_map: dict[str, Topic] = {} for name in self._store.list_streams(): - meta = self._store.get_stream_meta(name) - if meta: - channel = meta.get("channel") - if channel: - topic_map[name] = Topic.from_channel_str(channel) + s: Stream[object] = self._store.stream(name) + data_type = s.data_type + if data_type is not object: + topic_map[name] = Topic(topic=f"/{name}", lcm_type=data_type) # Merge-sort all streams by timestamp start_ts = t_min + self._position @@ -265,8 +259,7 @@ async def _playback_loop(self) -> None: counter = 0 # tiebreaker for heapq for name in self._store.list_streams(): - s: Stream[object] = self._store.stream(name) - it = iter(s.after(start_ts - 0.001)) + it: Any = iter(self._store.stream(name).after(start_ts - 0.001)) try: obs = next(it) except StopIteration: @@ -344,8 +337,7 @@ async def seek(self, seconds: float) -> None: self._position = max(0.0, min(seconds, self.duration)) if self.is_playing: await self.stop_playback() - assert self._pubsub is not None - self.play(pubsub=self._pubsub, speed=self._play_speed) + self.play(speed=self._play_speed) def delete_range(self, start: float, end: float) -> int: """Delete observations in [start, end] seconds from recording start. From 452130dc030a6234a9a574503afa452a44d78524 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 22:27:07 +0800 Subject: [PATCH 29/33] webrtc retransmit issue fix --- dimos/robot/unitree/connection.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index 85dd01ef55..103ac0ae52 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -17,7 +17,7 @@ import functools import threading import time -from typing import Any, TypeAlias +from typing import Any, TypeAlias, cast import numpy as np from numpy.typing import NDArray @@ -43,7 +43,7 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.type.lidar import RawLidarMsg, pointcloud2_from_webrtc_lidar from dimos.robot.unitree.type.lowstate import LowStateMsg -from dimos.robot.unitree.type.odometry import Odometry +from dimos.robot.unitree.type.odometry import Odometry, RawOdometryMessage from dimos.utils.decorators.decorators import simple_mcache from dimos.utils.reactive import backpressure, callback_to_observable @@ -240,11 +240,33 @@ def publish_request(self, topic: str, data: dict[Any, Any]) -> Any: @simple_mcache def raw_lidar_stream(self) -> Observable[RawLidarMsg]: - return backpressure(self.unitree_sub_stream(RTC_TOPIC["ULIDAR_ARRAY"])) + # Drop WebRTC retransmits at the source: under flaky-transport + # conditions Unitree's data channel occasionally redelivers the same + # logical frame, carrying an identical `data.stamp`. Both deliveries + # arrive back-to-back. Keeping both wastes downstream work (decode, + # codec, recording, viz) on a frame consumers already processed and + # leaves recordings with the same body_ts twice in a row. + return backpressure( + self.unitree_sub_stream(RTC_TOPIC["ULIDAR_ARRAY"]).pipe( + ops.distinct_until_changed( + key_mapper=lambda raw: cast("RawLidarMsg", raw)["data"]["stamp"] + ) + ) + ) @simple_mcache def raw_odom_stream(self) -> Observable[Pose]: - return backpressure(self.unitree_sub_stream(RTC_TOPIC["ROBOTODOM"])) + # Same WebRTC-retransmit dedup as raw_lidar_stream — odom just nests + # the stamp one level deeper inside `data.header`. + return backpressure( + self.unitree_sub_stream(RTC_TOPIC["ROBOTODOM"]).pipe( + ops.distinct_until_changed( + key_mapper=lambda raw: cast("RawOdometryMessage", raw)["data"]["header"][ + "stamp" + ] + ) + ) + ) @simple_mcache def lidar_stream(self) -> Observable[PointCloud2]: From 08337ce2a8fbfdee0134c300f542d748a4331086 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 22:53:59 +0800 Subject: [PATCH 30/33] webrtc lidar timestamps have been fixed --- dimos/robot/unitree/type/lidar.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/dimos/robot/unitree/type/lidar.py b/dimos/robot/unitree/type/lidar.py index 7df2f0ce0c..3d987470e0 100644 --- a/dimos/robot/unitree/type/lidar.py +++ b/dimos/robot/unitree/type/lidar.py @@ -14,7 +14,6 @@ """Unitree WebRTC lidar message parsing utilities.""" -import time from typing import TypedDict import numpy as np @@ -55,7 +54,8 @@ def pointcloud2_from_webrtc_lidar(raw_message: RawLidarMsg, ts: float | None = N Args: raw_message: Raw lidar message from Unitree WebRTC API - ts: Optional timestamp override. If None, uses current time. + ts: Optional timestamp override. If None, uses the sensor stamp from + ``raw_message["data"]["stamp"]``. Returns: PointCloud2 message with the lidar points @@ -68,7 +68,8 @@ def pointcloud2_from_webrtc_lidar(raw_message: RawLidarMsg, ts: float | None = N return PointCloud2( pointcloud=pointcloud, - # webrtc stamp is broken (e.g., "stamp": 1.758148e+09), use current time - ts=ts if ts is not None else time.time(), + # Use the sensor's own stamp — it's authoritative now that Unitree's + # WebRTC API reports a real unix timestamp here. + ts=ts if ts is not None else data["stamp"], frame_id="world", ) From edf2b76d4c8bdf67d0724c5bf9f1ebce53796153 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 22:54:32 +0800 Subject: [PATCH 31/33] record/replay should use wall clock for msg ts, not record time --- dimos/core/coordination/module_coordinator.py | 2 +- dimos/core/module.py | 7 ++++--- dimos/record/record_replay.py | 19 +++++-------------- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/dimos/core/coordination/module_coordinator.py b/dimos/core/coordination/module_coordinator.py index c71374afa3..b9285e583d 100644 --- a/dimos/core/coordination/module_coordinator.py +++ b/dimos/core/coordination/module_coordinator.py @@ -598,7 +598,7 @@ async def loop(self) -> None: stop = asyncio.Event() try: await stop.wait() - except KeyboardInterrupt: + except (KeyboardInterrupt, asyncio.CancelledError): return finally: self.stop() diff --git a/dimos/core/module.py b/dimos/core/module.py index d48ff63c33..8957d0fadb 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -46,7 +46,6 @@ from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable from dimos.protocol.tf.tf import LCMTF, TFSpec -from dimos.types.timestamped import Timestamped from dimos.utils import colors from dimos.utils.generic import classproperty from dimos.utils.logging_config import setup_logger @@ -426,8 +425,10 @@ def start_recording_outputs(self, db_path: str) -> None: stream = self._rec_store.stream(name, out.type) def cb(msg: Any, _stream: "Stream[object]" = stream) -> None: - ts = msg.ts if isinstance(msg, Timestamped) else time.time() - _stream.append(msg, ts=ts) + # Storage ts is wall-clock at record time. Sensor / message + # timestamps live on the payload (e.g. ``msg.ts``) and are + # preserved there. Storage ts must be unique within a stream. + _stream.append(msg, ts=time.time()) self._rec_unsubs.append(out.subscribe(cb)) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index e3e1b021eb..b65b22965f 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -53,7 +53,6 @@ from dimos.protocol.pubsub.impl.lcmpubsub import Topic from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.service.spec import Service -from dimos.types.timestamped import Timestamped from dimos.utils.logging_config import setup_logger from dimos.visualization.rerun.bridge import RerunConvertible, is_rerun_multi @@ -101,7 +100,6 @@ def __init__(self, path: str) -> None: self._recording = False self._unsubscribes: list[Callable[[], None]] = [] - self._ts_offsets: dict[str, float] = {} self._resume = asyncio.Event() self._play_task: asyncio.Task[None] | None = None @@ -160,19 +158,12 @@ def stop_recording(self) -> None: logger.info("Recording stopped") def _on_message(self, msg: object, topic: Topic) -> None: + # Storage ts = reception time (default in stream.append). Sensor stamp + # is preserved on the payload (msg.ts) — same split rosbag uses, and + # the only one that survives sender-side stamp bugs / retransmits / + # looped replay sources without UNIQUE collisions. stream_name = topic_to_stream_name(topic.pattern) - - ts: float | None = None - if isinstance(msg, Timestamped) and msg.ts: - # Record the time offset on first message and then offset all messages - # by the same amount. This allows recording live channels and replayed - # channels together with consistent timestamps. - if stream_name not in self._ts_offsets: - self._ts_offsets[stream_name] = time.time() - msg.ts - ts = msg.ts + self._ts_offsets[stream_name] - - s = self._store.stream(stream_name, type(msg)) - s.append(msg, ts=ts) + self._store.stream(stream_name, type(msg)).append(msg) @property def duration(self) -> float: From d66a7bbca0b639f07e5183daf58632b12cc63fd8 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 23:00:03 +0800 Subject: [PATCH 32/33] final cleanup --- dimos/record/record_replay.py | 1 - 1 file changed, 1 deletion(-) diff --git a/dimos/record/record_replay.py b/dimos/record/record_replay.py index b65b22965f..ae0e518cd9 100644 --- a/dimos/record/record_replay.py +++ b/dimos/record/record_replay.py @@ -154,7 +154,6 @@ def stop_recording(self) -> None: for unsub in self._unsubscribes: unsub() self._unsubscribes.clear() - self._ts_offsets.clear() logger.info("Recording stopped") def _on_message(self, msg: object, topic: Topic) -> None: From 5434e17760ef9e1217348a3abb0e1861b83311c1 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Tue, 5 May 2026 23:05:10 +0800 Subject: [PATCH 33/33] sql transaction fix --- dimos/memory2/observationstore/sqlite.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/dimos/memory2/observationstore/sqlite.py b/dimos/memory2/observationstore/sqlite.py index 34788e94c3..7189b61e06 100644 --- a/dimos/memory2/observationstore/sqlite.py +++ b/dimos/memory2/observationstore/sqlite.py @@ -448,17 +448,21 @@ def fetch_by_ids(self, ids: list[int]) -> list[Observation[T]]: def delete_range(self, t1: float, t2: float) -> list[int]: """Delete observations with ts in [t1, t2]. Returns deleted IDs.""" with self._lock: - rows = self._conn.execute( - f'DELETE FROM "{self._name}" WHERE ts >= ? AND ts <= ? RETURNING id', - (t1, t2), - ).fetchall() - ids = [r[0] for r in rows] - if ids: - placeholders = ",".join("?" * len(ids)) - self._conn.execute( - f'DELETE FROM "{self._name}_rtree" WHERE id IN ({placeholders})', ids - ) + try: + rows = self._conn.execute( + f'DELETE FROM "{self._name}" WHERE ts >= ? AND ts <= ? RETURNING id', + (t1, t2), + ).fetchall() + ids = [r[0] for r in rows] + if ids: + placeholders = ",".join("?" * len(ids)) + self._conn.execute( + f'DELETE FROM "{self._name}_rtree" WHERE id IN ({placeholders})', ids + ) self._conn.commit() + except Exception: + self._conn.rollback() + raise return ids def stop(self) -> None: