feat(unitree): forward aes_128_key to WebRTC driver for G1 firmware >=1.5.1

[mihai-chiorean]

Motivation

Unitree G1 firmware 1.5.1 introduced a data2=3 WebRTC handshake that requires a per-device AES-128 key. Without it the connection fails inside unitree_webrtc_connect with RSA key format is not supported (pycryptodome trying to parse ciphertext bytes as an RSA key).

The per-device key is fetched from Unitree's cloud once via unitree-fetch-aes-key --email YOU --sn <serial> and then cached locally. The upstream driver legion1581/unitree_webrtc_connect@v2.1.1 accepts it through an aes_128_key= kwarg on UnitreeWebRTCConnection. The dimos wrapper at dimos/robot/unitree/connection.py never forwarded it, so anyone on G1 firmware >=1.5.1 cannot use the dimos Unitree integration today.

Change

UnitreeWebRTCConnection.__init__ now accepts an optional aes_128_key: str | None = None kwarg and falls back to the UNITREE_AES_128_KEY environment variable. The value is forwarded to the underlying LegionConnection only when non-empty, so the call is byte-identical to the previous behaviour when unset.

def __init__(self, ip: str, mode: str = "ai", aes_128_key: str | None = None) -> None:
    ...
    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)

Verification

Tested on a Unitree G1 EDU+ (model string G1_Edu+) at firmware 1.5.1.1. With UNITREE_AES_128_KEY exported in the shell, dimos run unitree-g1-basic proceeds past the WebRTC handshake on 192.168.123.161:9991 and the robot accepts motion-switcher / wireless-controller commands as before. Without the env var, the same command reproduces RSA key format is not supported and the handshake never completes — same behaviour as main.

No breaking change

The kwarg is optional and defaults to None. The env-var lookup is opt-in (only consulted when the kwarg is None). Callers on G1 firmware <1.5.1 and all Go2 robots see no behavioural change — the **extra is empty and the LegionConnection call is byte-identical to the old one.

Optional follow-up (not in this PR)

pyproject.toml currently pins unitree-webrtc-connect-leshy>=2.0.7. The aes_128_key kwarg + the _decrypt_data1_v3 path that consumes it landed in legion1581/unitree_webrtc_connect@v2.1.1. If unitree-webrtc-connect-leshy is a downstream rebuild of that repo and is already on >=2.1.1, no change is needed and this PR is enough on its own. If it's still tracking 2.0.x, the dep will also need to be bumped (e.g. via [tool.uv.sources] pointing at git+https://github.com/legion1581/unitree_webrtc_connect.git@v2.1.1). Happy to follow up in a separate PR if maintainers confirm the situation.

[greptile-apps]

Greptile Summary

This PR wires an optional per-device AES-128 key through UnitreeWebRTCConnection so that G1 robots running firmware ≥ 1.5.1 can complete the data2=3 WebRTC handshake that was previously failing with RSA key format is not supported.

• The new aes_128_key kwarg is optional and falls back to an environment variable lookup, keeping the call to LegionConnection byte-identical to the old behaviour when the key is absent — no breaking change for older firmware or Go2 robots.
• The two higher-level G1 entry points (G1Connection and G1HighLevelWebRtc) do not yet expose aes_128_key in their module configs, so users of those paths must supply the key via the environment variable rather than declarative config.

Confidence Score: 5/5

Safe to merge; the change is additive, fully backward-compatible, and correctly guarded so existing callers see no behavioural difference.

The diff is small and self-contained: one new optional parameter with a truthiness guard, an env-var fallback, and a **extra splat that produces an empty dict when the key is absent. The only follow-up worth noting is that the two higher-level G1 module configs don't yet expose aes_128_key declaratively, but that is an ergonomic gap, not a defect — the env-var path works for all existing call sites.

No files require special attention; dimos/robot/unitree/g1/connection.py and dimos/robot/unitree/g1/effectors/high_level/webrtc.py could optionally grow aes_128_key config fields as a follow-up.

Important Files Changed

Filename	Overview
dimos/robot/unitree/connection.py	Adds optional aes_128_key kwarg to UnitreeWebRTCConnection.__init__ with env-var fallback; change is backward-compatible and correctly guards the kwarg with a truthiness check before passing it to LegionConnection.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[UnitreeWebRTCConnection.__init__ called] --> B{aes_128_key kwarg provided?}
    B -- Yes --> D{key non-empty?}
    B -- No --> C[Read key from environment]
    C --> D
    D -- Yes --> E[extra = aes_128_key kwarg dict]
    D -- No --> F[extra = empty dict]
    E --> G[LegionConnection with key forwarded]
    F --> H[LegionConnection unchanged call]
    G --> I[self.connect]
    H --> I

Loading

_{Reviews (1): Last reviewed commit: "feat(unitree): forward aes_128_key to We..." | Re-trigger Greptile}

[leshy]

reason for legion client fork is that unitree-webrtc-connect-leshy had a more efficinet lidar data parser and we had some issues with aiortc package pinning and had to fork it as well

legion1581/unitree_webrtc_connect@master...leshy:unitree_webrtc_connect:master
I assume lidar stuff is merged in the legion lib and all is resolved now

for FDEs (or anyone else here in the thread:) - we need to validate this on current g1s/go2s (ensure my lidar changes are now merged to legion client, transition to legion lib, test on our robots, make sure installation works, merge this)

[natcl]

We could also include GO2 firmware 1.1.15 as it has the same issue and this should fix it for the GO2 too !

[codecov]

Codecov Report

❌ Patch coverage is 33.33333% with 4 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
dimos/robot/unitree/connection.py	33.33%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!