feat: use individual module connections

.pre-commit-config.yaml

@@ -117,7 +117,7 @@ repos:
             git ls-files "*Cargo.toml" | while read -r m; do
               cargo locate-project --manifest-path "$m" --workspace --message-format plain
             done | sort -u | while read -r root; do
-              cargo fmt --manifest-path "$root" --all
+              cargo fmt --manifest-path "$root" --all --check
             done
           '
         language: system

dimos/agents/skills/person_follow.py

@@ -81,9 +81,9 @@ def __init__(self, **kwargs: Any) -> None:
         # Use simulator camera intrinsics in simulation mode
         camera_info = self.config.camera_info
         if self.config.g.simulation == "mujoco":
-            from dimos.robot.unitree.mujoco_connection import MujocoConnection
+            from dimos.robot.unitree.mujoco_camera_constants import MUJOCO_CAMERA_INFO_STATIC
 
-            camera_info = MujocoConnection.camera_info_static
+            camera_info = MUJOCO_CAMERA_INFO_STATIC
         elif self.config.g.simulation == "dimsim":
             from dimos.robot.unitree.dimsim_connection import DimSimConnection
 

dimos/core/coordination/blueprints.py

@@ -195,6 +195,66 @@ def requirements(self, *checks: Callable[[], str | None]) -> "Blueprint":
     def configurators(self, *checks: "SystemConfigurator") -> "Blueprint":
         return replace(self, configurator_checks=self.configurator_checks + tuple(checks))
 
+    def with_backend(self, backend: str) -> "Blueprint":
+        """Swap tagged connection modules for the matching `(robot, backend)` variant.
+
+        For each atom whose module carries a `_connection_tag` with a backend
+        different from the requested one, look up the same-robot module for
+        `backend` in the connection registry and substitute it. Streams,
+        remappings, and disabled-modules entries are rewritten to point at the
+        new class.
+
+        If the blueprint has no tagged atoms this is a no-op (with a warning).
+        """
+        # Lazy import to keep blueprints.py free of robot deps.
+        from dimos.robot.connection_registry import backends_for, get_connection
+
+        swap_map: dict[type[ModuleBase], type[ModuleBase]] = {}
+        for atom in self.blueprints:
+            tag = getattr(atom.module, "_connection_tag", None)
+            if tag is None or tag.backend == backend:
+                continue
+            target = get_connection(tag.robot, backend)
+            if target is None:
+                available = sorted(backends_for(tag.robot))
+                raise ValueError(
+                    f"No connection registered for robot={tag.robot!r} "
+                    f"backend={backend!r} (have: {available})"
+                )
+            swap_map[atom.module] = target
+
+        if not swap_map:
+            tagged = any(getattr(a.module, "_connection_tag", None) for a in self.blueprints)
+            if not tagged:
+                logger.warning(
+                    "Blueprint.with_backend(%r) had no tagged connection atoms — "
+                    "returning blueprint unchanged",
+                    backend,
+                )
+            return self
+
+        new_atoms: list[BlueprintAtom] = []
+        for atom in self.blueprints:
+            target = swap_map.get(atom.module)
+            if target is None:
+                new_atoms.append(atom)
+                continue
+            _check_stream_parity(atom.module, target, atom)
+            _check_kwargs_compat(target, atom.kwargs)
+            new_atoms.append(BlueprintAtom.create(target, atom.kwargs))
+
+        new_remappings = {
+            (swap_map.get(m, m), name): v for (m, name), v in self.remapping_map.items()
+        }
+        new_disabled = tuple(swap_map.get(m, m) for m in self.disabled_modules_tuple)
+
+        return replace(
+            self,
+            blueprints=tuple(new_atoms),
+            remapping_map=MappingProxyType(new_remappings),
+            disabled_modules_tuple=new_disabled,
+        )
+
     @cached_property
     def active_blueprints(self) -> tuple[BlueprintAtom, ...]:
         if not self.disabled_modules_tuple:
@@ -239,3 +299,38 @@ def _eliminate_duplicates(blueprints: list[BlueprintAtom]) -> list[BlueprintAtom
             seen.add(bp.module)
             unique_blueprints.append(bp)
     return list(reversed(unique_blueprints))
+
+
+def _stream_signature(streams: tuple[StreamRef, ...]) -> set[tuple[str, str]]:
+    return {(s.name, s.direction) for s in streams}
+
+
+def _check_stream_parity(old: type[ModuleBase], new: type[ModuleBase], atom: BlueprintAtom) -> None:
+    new_atom = BlueprintAtom.create(new, atom.kwargs)
+    old_sig = _stream_signature(atom.streams)
+    new_sig = _stream_signature(new_atom.streams)
+    if old_sig != new_sig:
+        only_old = sorted(old_sig - new_sig)
+        only_new = sorted(new_sig - old_sig)
+        raise ValueError(
+            f"Stream surface drift swapping {old.__name__} -> {new.__name__}: "
+            f"only on {old.__name__}={only_old}, only on {new.__name__}={only_new}"
+        )
+
+
+def _check_kwargs_compat(new: type[ModuleBase], kwargs: dict[str, Any]) -> None:
+    if not kwargs:
+        return
+    try:
+        config_type = get_type_hints(new).get("config")
+    except Exception:
+        return
+    if config_type is None:
+        return
+    valid_fields = set(getattr(config_type, "model_fields", {}))
+    invalid = set(kwargs) - valid_fields
+    if invalid:
+        raise ValueError(
+            f"Kwargs from blueprint atom are incompatible with {new.__name__}'s "
+            f"config ({config_type.__name__}): unknown field(s) {sorted(invalid)}"
+        )

dimos/core/coordination/test_blueprints.py

@@ -33,9 +33,11 @@
     autoconnect,
 )
 from dimos.core.core import rpc
-from dimos.core.module import Module
+from dimos.core.module import Module, ModuleConfig
 from dimos.core.stream import In, Out
 from dimos.core.transport import LCMTransport
+from dimos.robot import connection_registry
+from dimos.robot.connection_registry import connection
 from dimos.spec.utils import Spec
 
 
@@ -232,3 +234,139 @@ def test_active_blueprints_filters_disabled() -> None:
     active_modules = {bp.module for bp in blueprint.active_blueprints}
     assert ModuleA not in active_modules
     assert ModuleB in active_modules
+
+
+@pytest.fixture
+def isolated_registry(monkeypatch):
+    monkeypatch.setattr(connection_registry, "_REGISTRY", {})
+    yield connection_registry._REGISTRY
+
+
+class _BotConfig(ModuleConfig):
+    setting: str = "default"
+
+
+def _bot_modules():
+    """Create three (robot=bot) connection variants in an isolated registry."""
+
+    @connection(robot="bot", backend="real")
+    class BotReal(Module):
+        config: _BotConfig
+        cmd: In[Data1]
+        odom: Out[Data2]
+
+    @connection(robot="bot", backend="sim")
+    class BotSim(Module):
+        config: _BotConfig
+        cmd: In[Data1]
+        odom: Out[Data2]
+
+    @connection(robot="bot", backend="replay")
+    class BotReplay(Module):
+        config: ModuleConfig
+        cmd: In[Data1]
+        odom: Out[Data2]
+
+    return BotReal, BotSim, BotReplay
+
+
+def test_with_backend_no_op_when_no_tagged_atoms(isolated_registry) -> None:
+    blueprint = autoconnect(ModuleA.blueprint(), ModuleB.blueprint())
+    swapped = blueprint.with_backend("sim")
+    assert swapped is blueprint
+
+
+def test_with_backend_swaps_tagged_atom(isolated_registry) -> None:
+    BotReal, BotSim, _ = _bot_modules()
+
+    blueprint = autoconnect(ModuleA.blueprint(), BotReal.blueprint(setting="x"))
+    swapped = blueprint.with_backend("sim")
+
+    swapped_modules = [a.module for a in swapped.blueprints]
+    assert BotSim in swapped_modules
+    assert BotReal not in swapped_modules
+    assert ModuleA in swapped_modules
+
+    bot_atom = next(a for a in swapped.blueprints if a.module is BotSim)
+    assert bot_atom.kwargs == {"setting": "x"}
+    # Streams were re-extracted from the new class.
+    assert {s.name for s in bot_atom.streams} == {"cmd", "odom"}
+
+
+def test_with_backend_no_op_when_already_target(isolated_registry) -> None:
+    BotReal, _, _ = _bot_modules()
+
+    blueprint = BotReal.blueprint()
+    swapped = blueprint.with_backend("real")
+    assert swapped is blueprint  # no atoms needed swapping; returns self
+
+
+def test_with_backend_unknown_backend_raises(isolated_registry) -> None:
+    BotReal, _, _ = _bot_modules()
+
+    blueprint = BotReal.blueprint()
+    with pytest.raises(ValueError, match="No connection registered.*backend='nope'"):
+        blueprint.with_backend("nope")
+
+
+def test_with_backend_rewrites_remappings(isolated_registry) -> None:
+    BotReal, BotSim, _ = _bot_modules()
+
+    blueprint = BotReal.blueprint().remappings([(BotReal, "cmd", "remapped_cmd")])
+    swapped = blueprint.with_backend("sim")
+
+    assert (BotReal, "cmd") not in swapped.remapping_map
+    assert swapped.remapping_map[(BotSim, "cmd")] == "remapped_cmd"
+
+
+def test_with_backend_rewrites_disabled_modules(isolated_registry) -> None:
+    BotReal, BotSim, _ = _bot_modules()
+
+    blueprint = autoconnect(BotReal.blueprint(), ModuleA.blueprint()).disabled_modules(BotReal)
+    swapped = blueprint.with_backend("sim")
+
+    assert BotReal not in swapped.disabled_modules_tuple
+    assert BotSim in swapped.disabled_modules_tuple
+
+
+def test_with_backend_eager_kwarg_validation_raises(isolated_registry) -> None:
+    @connection(robot="bot", backend="real")
+    class BotReal2(Module):
+        class Cfg(ModuleConfig):
+            mode: str = "default"
+            speed: int = 1
+
+        config: Cfg
+        cmd: In[Data1]
+
+    @connection(robot="bot", backend="sim")
+    class BotSim2(Module):
+        class Cfg(ModuleConfig):
+            speed: int = 1  # NOTE: no `mode` field
+
+        config: Cfg
+        cmd: In[Data1]
+
+    blueprint = BotReal2.blueprint(mode="rage")
+    with pytest.raises(ValueError, match="unknown field.*mode"):
+        blueprint.with_backend("sim")
+
+
+def test_with_backend_stream_parity_drift_raises(isolated_registry) -> None:
+    @connection(robot="bot", backend="real")
+    class BotReal3(Module):
+        config: ModuleConfig
+        cmd: In[Data1]
+        odom: Out[Data2]
+
+    @connection(robot="bot", backend="sim")
+    class BotSim3(Module):
+        config: ModuleConfig
+        cmd: In[Data1]
+        # missing odom; adds extra stream
+
+        extra: Out[Data3]
+
+    blueprint = BotReal3.blueprint()
+    with pytest.raises(ValueError, match="Stream surface drift"):
+        blueprint.with_backend("sim")

dimos/core/global_config.py

@@ -89,6 +89,18 @@ def unitree_connection_type(self) -> str:
             return self.simulation
         return "webrtc"
 
+    @property
+    def effective_simulator(self) -> str | None:
+        """Connection backend to swap blueprint atoms to via `with_backend`.
+
+        `--replay` and `--simulation <name>` each map onto a registered
+        `(robot, backend)` connection. Returns None for the real-hardware
+        default (webrtc), which leaves the blueprint's atoms unchanged.
+        """
+        if self.replay:
+            return "replay"
+        return self.simulation or None
+
     @property
     def mujoco_start_pos_float(self) -> tuple[float, float]:
         x, y = _get_all_numbers(self.mujoco_start_pos)

dimos/core/module.py

@@ -105,6 +105,12 @@ class _BlueprintPartial(Protocol):
     def __call__(self, **kwargs: Any) -> "Blueprint": ...
 
 
+@dataclass(frozen=True)
+class ConnectionTag:
+    robot: str
+    backend: str
+
+
 class ModuleBase(Configurable, CompositeResource):
     config: ModuleConfig
 
@@ -128,6 +134,7 @@ class ModuleBase(Configurable, CompositeResource):
     _main_gen: AsyncGenerator[None, None] | None = None
     _tools: dict[str, Any]
     _tools_lock: threading.Lock
+    _connection_tag: ConnectionTag | None = None
 
     def __init__(self, config_args: dict[str, Any]) -> None:
         super().__init__(**config_args)

dimos/core/test_global_config.py

@@ -12,16 +12,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for GlobalConfig security defaults."""
+from dimos.core.global_config import GlobalConfig
 
 
 class TestGlobalConfigSecurityDefaults:
     """Network services must bind to localhost by default (not 0.0.0.0)."""
 
     def test_listen_host_defaults_to_localhost(self) -> None:
-        from dimos.core.global_config import GlobalConfig
-
         config = GlobalConfig()
         assert config.listen_host == "127.0.0.1", (
             f"listen_host must default to 127.0.0.1, got {config.listen_host}"
         )
+
+
+class TestSimulatorBackendResolution:
+    """`--simulation <name>` selects the connection backend; `--replay` wins."""
+
+    def test_simulation_selects_named_backend(self) -> None:
+        config = GlobalConfig(simulation="dimsim")
+        assert config.effective_simulator == "dimsim"
+        assert config.unitree_connection_type == "dimsim"
+
+    def test_simulation_mujoco(self) -> None:
+        config = GlobalConfig(simulation="mujoco")
+        assert config.effective_simulator == "mujoco"
+        assert config.unitree_connection_type == "mujoco"
+
+    def test_unset_returns_none_and_webrtc(self) -> None:
+        config = GlobalConfig(simulation="")
+        assert config.effective_simulator is None
+        assert config.unitree_connection_type == "webrtc"
+
+    def test_replay_overrides_simulation(self) -> None:
+        config = GlobalConfig(replay=True, simulation="mujoco")
+        assert config.effective_simulator == "replay"
+        assert config.unitree_connection_type == "replay"

dimos/experimental/security_demo/security_module.py

@@ -119,9 +119,9 @@ def _create_visual_servo(
 ) -> VisualServoing2D:
     camera_info = config.camera_info
     if global_config.simulation:
-        from dimos.robot.unitree.mujoco_connection import MujocoConnection
+        from dimos.robot.unitree.mujoco_camera_constants import MUJOCO_CAMERA_INFO_STATIC
 
-        camera_info = MujocoConnection.camera_info_static
+        camera_info = MUJOCO_CAMERA_INFO_STATIC
 
     return VisualServoing2D(camera_info, bool(global_config.simulation))
 

dimos/memory2/test_visualizer.py

@@ -29,7 +29,7 @@
 from dimos.msgs.geometry_msgs.Transform import Transform
 from dimos.msgs.sensor_msgs.Image import Image
 from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC
-from dimos.robot.unitree.go2.connection import GO2Connection
+from dimos.robot.unitree.go2.camera import GO2_CAMERA_INFO_STATIC
 from dimos.utils.data import get_data, get_data_dir
 
 if TYPE_CHECKING:
@@ -192,7 +192,7 @@ def test_detect_objects_smart(self, store: SqliteStore, clip: CLIPModel) -> None
                 det3d = Detection3DPC.from_2d(
                     det,
                     lidar_frame,
-                    camera_info=GO2Connection.camera_info_static,
+                    camera_info=GO2_CAMERA_INFO_STATIC,
                     world_to_optical_transform=Transform(
                         ts=obs.ts,
                         translation=obs.pose_stamped.position,