From f39ee53dfa1984e1568638ef80eca4424a3202be Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 16 May 2026 15:48:48 -0700 Subject: [PATCH 1/4] feat(unitree): forward aes_128_key to WebRTC driver for G1 firmware >=1.5.1 Unitree G1 firmware 1.5.1 added a data2=3 WebRTC handshake that requires a per-device AES-128 key (fetched from Unitree cloud once via `unitree-fetch-aes-key --email YOU --sn `). The upstream `legion1581/unitree_webrtc_connect@v2.1.1` accepts this via the `aes_128_key=` kwarg; without it the handshake fails with `RSA key format is not supported` from pycryptodome reading ciphertext bytes. `UnitreeWebRTCConnection.__init__` now accepts an optional `aes_128_key` kwarg and falls back to the `UNITREE_AES_128_KEY` environment variable so existing call sites do not need to change. When the value is unset the call to `LegionConnection` is byte-identical to the previous behaviour, so this is a no-op for G1 firmware <1.5.1 and for Go2. Verified against a Unitree G1 EDU+ on firmware 1.5.1.1 via `dimos run unitree-g1-basic` - the WebRTC handshake on 192.168.123.161:9991 succeeds with UNITREE_AES_128_KEY exported and reproduces the `RSA key format is not supported` failure without it. --- dimos/robot/unitree/connection.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index 919efc76f6..f6f48f27aa 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -15,6 +15,7 @@ import asyncio from dataclasses import dataclass import functools +import os import threading import time from typing import Any, TypeAlias @@ -85,12 +86,17 @@ def to_ndarray(self, format=None): # type: ignore[no-untyped-def] class UnitreeWebRTCConnection(Resource): _SPORT_API_ID_RAGEMODE: int = 2059 - def __init__(self, ip: str, mode: str = "ai") -> None: + def __init__(self, ip: str, mode: str = "ai", aes_128_key: str | None = None) -> None: self.ip = ip self.mode = mode self.stop_timer: threading.Timer | None = None self.cmd_vel_timeout = 0.2 - self.conn = LegionConnection(WebRTCConnectionMethod.LocalSTA, ip=self.ip) + # Per-device AES-128 key required by G1 firmware >= 1.5.1 (data2=3 WebRTC handshake). + # Fetch with: unitree-fetch-aes-key --email YOU --sn + if aes_128_key is None: + aes_128_key = os.environ.get("UNITREE_AES_128_KEY") + extra: dict[str, Any] = {"aes_128_key": aes_128_key} if aes_128_key else {} + self.conn = LegionConnection(WebRTCConnectionMethod.LocalSTA, ip=self.ip, **extra) self.connect() def connect(self) -> None: From a98a5f86187506e06f378349972f80af6264bd83 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sun, 17 May 2026 11:28:01 -0700 Subject: [PATCH 2/4] feat(unitree): expose aes_128_key in G1 module configs + add unit tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR review feedback from #2117: 1. G1Config and G1HighLevelWebRtcConfig now expose 'aes_128_key' as an optional declarative config field, forwarded to UnitreeWebRTCConnection. Users on G1 firmware >= 1.5.1 can now set it via blueprint config without resorting to the UNITREE_AES_128_KEY env var. The env-var fallback in UnitreeWebRTCConnection.__init__ is preserved as-is. 2. Adds dimos/robot/unitree/test_connection.py with five test cases that cover the kwarg-forwarding logic without hardware: - default behaviour: no kwarg, no env → aes_128_key not forwarded (byte-identical to pre-PR call) - explicit kwarg is forwarded verbatim - env-var fallback when kwarg is None - explicit kwarg beats env-var (precedence) - empty-string kwarg is treated as 'no key' (truthiness guard) Tests are pure-Python: LegionConnection is mocked, .connect() is patched to a no-op. Default pytest selector (-m 'not tool ...') will run them. --- dimos/robot/unitree/g1/connection.py | 7 +- .../unitree/g1/effectors/high_level/webrtc.py | 9 +- dimos/robot/unitree/test_connection.py | 91 +++++++++++++++++++ 3 files changed, 105 insertions(+), 2 deletions(-) create mode 100644 dimos/robot/unitree/test_connection.py diff --git a/dimos/robot/unitree/g1/connection.py b/dimos/robot/unitree/g1/connection.py index 58b2a0747f..253d76eda7 100644 --- a/dimos/robot/unitree/g1/connection.py +++ b/dimos/robot/unitree/g1/connection.py @@ -36,6 +36,9 @@ class G1Config(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) connection_type: str = Field(default_factory=lambda m: m["g"].unitree_connection_type) + # Per-device AES-128 key for G1 firmware >= 1.5.1 (data2=3 WebRTC handshake). + # If unset here, UnitreeWebRTCConnection falls back to the UNITREE_AES_128_KEY env var. + aes_128_key: str | None = None class G1ConnectionBase(Module, ABC): @@ -78,7 +81,9 @@ def start(self) -> None: match self.config.connection_type: case "webrtc": - self.connection = UnitreeWebRTCConnection(self.config.ip) + self.connection = UnitreeWebRTCConnection( + self.config.ip, aes_128_key=self.config.aes_128_key + ) case "replay": raise ValueError("Replay connection not implemented for G1 robot") case "mujoco": diff --git a/dimos/robot/unitree/g1/effectors/high_level/webrtc.py b/dimos/robot/unitree/g1/effectors/high_level/webrtc.py index 1784605ac2..134cd5c2d7 100644 --- a/dimos/robot/unitree/g1/effectors/high_level/webrtc.py +++ b/dimos/robot/unitree/g1/effectors/high_level/webrtc.py @@ -44,6 +44,9 @@ class G1HighLevelWebRtcConfig(ModuleConfig): ip: str | None = None connection_mode: str = "ai" + # Per-device AES-128 key for G1 firmware >= 1.5.1 (data2=3 WebRTC handshake). + # If unset here, UnitreeWebRTCConnection falls back to the UNITREE_AES_128_KEY env var. + aes_128_key: str | None = None class G1HighLevelWebRtc(Module, HighLevelG1Spec): @@ -62,7 +65,11 @@ def __init__(self, *args: Any, g: GlobalConfig = global_config, **kwargs: Any) - def start(self) -> None: super().start() assert self.config.ip is not None, "ip must be set in G1HighLevelWebRtcConfig" - self.connection = UnitreeWebRTCConnection(self.config.ip, self.config.connection_mode) + self.connection = UnitreeWebRTCConnection( + self.config.ip, + self.config.connection_mode, + aes_128_key=self.config.aes_128_key, + ) self.connection.start() self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) diff --git a/dimos/robot/unitree/test_connection.py b/dimos/robot/unitree/test_connection.py new file mode 100644 index 0000000000..60c92ad069 --- /dev/null +++ b/dimos/robot/unitree/test_connection.py @@ -0,0 +1,91 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for UnitreeWebRTCConnection's aes_128_key kwarg + env-var fallback. + +Pure-Python tests — no hardware, no network. Mocks the LegionConnection driver +and the connect() side-effect so __init__ stays inside the kwarg-forwarding logic. +""" + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from dimos.robot.unitree import connection as conn_mod +from dimos.robot.unitree.connection import UnitreeWebRTCConnection + + +@pytest.fixture +def stub_legion(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replace LegionConnection in the module with a MagicMock and suppress + UnitreeWebRTCConnection.connect so __init__ doesn't try to dial out.""" + monkeypatch.setattr(UnitreeWebRTCConnection, "connect", lambda self: None) + legion = MagicMock(name="LegionConnection") + monkeypatch.setattr(conn_mod, "LegionConnection", legion) + return legion + + +def _aes_kwarg(legion: MagicMock) -> Any: + """Pull aes_128_key out of the LegionConnection call args, or None if absent.""" + _args, kwargs = legion.call_args + return kwargs.get("aes_128_key") + + +def test_aes_key_omitted_when_neither_kwarg_nor_env( + stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + """Default behaviour: no kwarg, no env var → aes_128_key not forwarded. + + Guarantees the call is byte-identical to the pre-PR behaviour for users + on G1 firmware <1.5.1 and all Go2 robots. + """ + monkeypatch.delenv("UNITREE_AES_128_KEY", raising=False) + UnitreeWebRTCConnection(ip="192.168.123.161") + assert _aes_kwarg(stub_legion) is None + assert "aes_128_key" not in stub_legion.call_args.kwargs + + +def test_aes_key_from_explicit_kwarg( + stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + """Caller passes the key directly → forwarded verbatim.""" + monkeypatch.delenv("UNITREE_AES_128_KEY", raising=False) + UnitreeWebRTCConnection(ip="192.168.123.161", aes_128_key="aa" * 16) + assert _aes_kwarg(stub_legion) == "aa" * 16 + + +def test_aes_key_from_env_when_kwarg_none( + stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + """Env-var fallback: kwarg unset → UNITREE_AES_128_KEY is used.""" + monkeypatch.setenv("UNITREE_AES_128_KEY", "bb" * 16) + UnitreeWebRTCConnection(ip="192.168.123.161") + assert _aes_kwarg(stub_legion) == "bb" * 16 + + +def test_explicit_kwarg_beats_env(stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + """Precedence: explicit kwarg wins over UNITREE_AES_128_KEY env var.""" + monkeypatch.setenv("UNITREE_AES_128_KEY", "from-env") + UnitreeWebRTCConnection(ip="192.168.123.161", aes_128_key="from-kwarg") + assert _aes_kwarg(stub_legion) == "from-kwarg" + + +def test_empty_string_kwarg_skips_forwarding( + stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + """An empty-string kwarg is treated as 'no key' (truthiness guard).""" + monkeypatch.delenv("UNITREE_AES_128_KEY", raising=False) + UnitreeWebRTCConnection(ip="192.168.123.161", aes_128_key="") + assert "aes_128_key" not in stub_legion.call_args.kwargs From fceccd7c004b665ca0780c430fd2f62190050dfa Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sun, 17 May 2026 11:33:03 -0700 Subject: [PATCH 3/4] Update dimos/robot/unitree/connection.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- dimos/robot/unitree/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index f6f48f27aa..03fec366d0 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -93,7 +93,7 @@ def __init__(self, ip: str, mode: str = "ai", aes_128_key: str | None = None) -> self.cmd_vel_timeout = 0.2 # Per-device AES-128 key required by G1 firmware >= 1.5.1 (data2=3 WebRTC handshake). # Fetch with: unitree-fetch-aes-key --email YOU --sn - if aes_128_key is None: + if not aes_128_key: aes_128_key = os.environ.get("UNITREE_AES_128_KEY") extra: dict[str, Any] = {"aes_128_key": aes_128_key} if aes_128_key else {} self.conn = LegionConnection(WebRTCConnectionMethod.LocalSTA, ip=self.ip, **extra) From 6bff60d448e9ad84fee882ae36b5991d2c780de6 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Mon, 18 May 2026 14:17:37 -0700 Subject: [PATCH 4/4] feat(unitree): thread aes_128_key through Go2 connection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the G1 wiring from this PR so Go2 firmware >= 1.1.15 (same data2=3 WebRTC handshake change) can pass the per-device AES-128 key. Adds the field to Go2's ConnectionConfig, accepts it as a kwarg on make_connection, and forwards it to UnitreeWebRTCConnection in the webrtc branch. Other branches (replay, mujoco, dimsim) are unaffected. Also fixes the test docstring greptile flagged after the 'if not aes_128_key' guard landed: the test asserting empty-string + unset env → no forwarding had a misleading docstring claiming the kwarg was 'treated as no key'. Rewrites the docstring and adds a sibling test covering the case that actually motivated the guard fix: empty-string kwarg + set env → env value wins (the YAML/JSON 'aes_128_key: ""' regression case). Addresses natcl's review comment on PR #2117. --- dimos/robot/unitree/go2/connection.py | 15 +++++-- dimos/robot/unitree/go2/fleet_connection.py | 7 ++- dimos/robot/unitree/go2/test_connection.py | 50 +++++++++++++++++++++ dimos/robot/unitree/test_connection.py | 26 ++++++++++- 4 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 dimos/robot/unitree/go2/test_connection.py diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 4083def93b..b677f58cd6 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -64,6 +64,9 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT + # Per-device AES-128 key for Go2 firmware >= 1.1.15 (data2=3 WebRTC handshake). + # If unset here, UnitreeWebRTCConnection falls back to the UNITREE_AES_128_KEY env var. + aes_128_key: str | None = None class Go2ConnectionProtocol(Protocol): @@ -113,7 +116,11 @@ def _camera_info_static() -> CameraInfo: ) -def make_connection(ip: str | None, cfg: GlobalConfig) -> Go2ConnectionProtocol: +def make_connection( + ip: str | None, + cfg: GlobalConfig, + aes_128_key: str | None = None, +) -> Go2ConnectionProtocol: connection_type = cfg.unitree_connection_type if ip in ("fake", "mock", "replay") or connection_type == "replay": @@ -129,7 +136,7 @@ def make_connection(ip: str | None, cfg: GlobalConfig) -> Go2ConnectionProtocol: return DimSimConnection(cfg) elif connection_type == "webrtc": assert ip is not None, "IP address must be provided" - return UnitreeWebRTCConnection(ip) + return UnitreeWebRTCConnection(ip, aes_128_key=aes_128_key) else: raise ValueError(f"Unknown simulator {cfg.simulation!r}. Choose from: mujoco, dimsim") @@ -221,7 +228,9 @@ def rerun_views(cls): # type: ignore[no-untyped-def] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - self.connection = make_connection(self.config.ip, self.config.g) + self.connection = make_connection( + self.config.ip, self.config.g, aes_128_key=self.config.aes_128_key + ) if hasattr(self.connection, "camera_info_static"): self.camera_info_static = self.connection.camera_info_static diff --git a/dimos/robot/unitree/go2/fleet_connection.py b/dimos/robot/unitree/go2/fleet_connection.py index f2a0216ab7..e17e66f542 100644 --- a/dimos/robot/unitree/go2/fleet_connection.py +++ b/dimos/robot/unitree/go2/fleet_connection.py @@ -71,7 +71,12 @@ def __init__(self, **kwargs: Any) -> None: def start(self) -> None: self._extra_connections.clear() for ip in self._extra_ips: - conn = make_connection(ip, self.config.g) + # Forward the configured key to every follower. The AES-128 key is + # technically per-device, so a heterogeneous fleet needs each robot's + # own key; that's a future per-IP mapping. For homogeneous-key or + # env-var-driven setups this matches the primary's behaviour and + # avoids leaving followers reliant on UNITREE_AES_128_KEY alone. + conn = make_connection(ip, self.config.g, aes_128_key=self.config.aes_128_key) conn.start() self._extra_connections.append(conn) diff --git a/dimos/robot/unitree/go2/test_connection.py b/dimos/robot/unitree/go2/test_connection.py new file mode 100644 index 0000000000..28445826b1 --- /dev/null +++ b/dimos/robot/unitree/go2/test_connection.py @@ -0,0 +1,50 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Targeted test for go2.connection.make_connection's aes_128_key forwarding. + +Pins the wiring added in PR #2117 so renaming the kwarg in either +UnitreeWebRTCConnection or make_connection without updating the other +fails loudly. The leaf (UnitreeWebRTCConnection.__init__) is exercised +in dimos/robot/unitree/test_connection.py — this file only covers the +go2-local routing. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from dimos.robot.unitree.go2 import connection as go2_conn + + +@pytest.fixture +def stub_webrtc(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replace UnitreeWebRTCConnection in go2.connection with a MagicMock so + make_connection's webrtc branch is exercised without dialing out.""" + stub = MagicMock(name="UnitreeWebRTCConnection") + monkeypatch.setattr(go2_conn, "UnitreeWebRTCConnection", stub) + return stub + + +def test_make_connection_webrtc_forwards_aes_128_key(stub_webrtc: MagicMock) -> None: + """Webrtc branch must forward aes_128_key as a kwarg to UnitreeWebRTCConnection. + + Guards the Go2 half of the PR #2117 fix: without this forwarding, Go2 + robots on firmware >=1.1.15 fail the WebRTC handshake even when a key + is provided in config. + """ + cfg = SimpleNamespace(unitree_connection_type="webrtc") + go2_conn.make_connection("192.168.123.161", cfg, aes_128_key="cafe" * 8) + stub_webrtc.assert_called_once_with("192.168.123.161", aes_128_key="cafe" * 8) diff --git a/dimos/robot/unitree/test_connection.py b/dimos/robot/unitree/test_connection.py index 60c92ad069..56f26479f8 100644 --- a/dimos/robot/unitree/test_connection.py +++ b/dimos/robot/unitree/test_connection.py @@ -82,10 +82,32 @@ def test_explicit_kwarg_beats_env(stub_legion: MagicMock, monkeypatch: pytest.Mo assert _aes_kwarg(stub_legion) == "from-kwarg" -def test_empty_string_kwarg_skips_forwarding( +def test_empty_string_kwarg_falls_back_to_env_when_unset( stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch ) -> None: - """An empty-string kwarg is treated as 'no key' (truthiness guard).""" + """Empty-string kwarg + unset env → final key is None, nothing forwarded. + + Truthiness guard means `aes_128_key=""` falls through to the env-var + lookup just like `aes_128_key=None` would. With the env unset, the final + value is None and the kwarg is omitted from the LegionConnection call — + same byte-identical behaviour as the unset case. + """ monkeypatch.delenv("UNITREE_AES_128_KEY", raising=False) UnitreeWebRTCConnection(ip="192.168.123.161", aes_128_key="") assert "aes_128_key" not in stub_legion.call_args.kwargs + + +def test_empty_string_kwarg_uses_env_when_set( + stub_legion: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + """Empty-string kwarg + set env → env value wins. + + Guards against the bug greptile flagged on PR #2117: a YAML/JSON config + serialising `aes_128_key: ""` instead of `null` MUST still allow the env + var to provide the real key. Otherwise the connection silently fails on + G1 firmware >=1.5.1 (and Go2 firmware >=1.1.15) with + 'RSA key format is not supported'. + """ + monkeypatch.setenv("UNITREE_AES_128_KEY", "cc" * 16) + UnitreeWebRTCConnection(ip="192.168.123.161", aes_128_key="") + assert _aes_kwarg(stub_legion) == "cc" * 16