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 threads an optional aes_128_key parameter through the Unitree WebRTC connection stack so that G1 robots on firmware ≥1.5.1 can complete the data2=3 WebRTC handshake that requires a per-device AES-128 key.

• UnitreeWebRTCConnection.__init__ now accepts aes_128_key: str | None = None, falls back to UNITREE_AES_128_KEY env var when the kwarg is falsy, and forwards it only when non-empty — keeping older firmware paths byte-identical.
• Both G1Config and G1HighLevelWebRtcConfig expose the same field so callers can supply the key via config objects.
• A new pure-Python test file (test_connection.py) mocks LegionConnection and covers the key forwarding scenarios.

Confidence Score: 5/5

Safe to merge — the change is additive and backward-compatible; all existing callers are unaffected when neither the kwarg nor the env var is set.

The core logic in connection.py correctly uses if not aes_128_key: (handling both None and "") before the env-var lookup, and only injects aes_128_key into LegionConnection when the value is non-empty. The config additions in G1Config and G1HighLevelWebRtcConfig are straightforward pass-throughs. The only gap is a test docstring that understates the env-var fallback behaviour for empty-string inputs, which does not affect correctness.

No files require special attention.

Important Files Changed

Filename	Overview
dimos/robot/unitree/connection.py	Core change: adds aes_128_key kwarg with env-var fallback and conditional forwarding to LegionConnection; logic is correct.
dimos/robot/unitree/g1/connection.py	Adds aes_128_key field to G1Config and threads it through to UnitreeWebRTCConnection; straightforward and correct.
dimos/robot/unitree/g1/effectors/high_level/webrtc.py	Adds aes_128_key field to G1HighLevelWebRtcConfig and passes it to UnitreeWebRTCConnection; mirrors the pattern in g1/connection.py consistently.
dimos/robot/unitree/test_connection.py	Good coverage of the happy paths; test_empty_string_kwarg_skips_forwarding has a misleading docstring — empty string actually triggers env-var lookup, not unconditional "no key".

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[UnitreeWebRTCConnection.__init__] --> B{aes_128_key kwarg truthy?}
    B -- Yes --> D[Build extra dict with key]
    B -- No --> C[Read AES key from environment]
    C --> E{Env key truthy?}
    E -- Yes --> D
    E -- No --> F[extra = empty dict]
    D --> G[LegionConnection with aes key]
    F --> H[LegionConnection without aes key]

Loading

_{Reviews (4): Last reviewed commit: "Update dimos/robot/unitree/connection.py" | 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!

[natcl]

had a more efficinet lidar data parser

I did notice a degradation in navigation performance on our go2 edu when I switched to the upstream library, might be worth checking if the changes are there.