From 160db3fecfb32883f84eae8525979fb7410dc043 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 7 Aug 2026 19:17:09 -0700 Subject: [PATCH 1/4] feat(teleop): add LiveKit hosted teleop module --- dimos/robot/all_blueprints.py | 6 + dimos/teleop/hosted/blueprints/livekit.py | 121 ++++++ dimos/teleop/hosted/livekit.py | 445 ++++++++++++++++++++++ dimos/teleop/hosted/test_livekit.py | 215 +++++++++++ pyproject.toml | 12 +- uv.lock | 59 ++- 6 files changed, 845 insertions(+), 13 deletions(-) create mode 100644 dimos/teleop/hosted/blueprints/livekit.py create mode 100644 dimos/teleop/hosted/livekit.py create mode 100644 dimos/teleop/hosted/test_livekit.py diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 58b941ae79..72d7fe6cd7 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -91,10 +91,13 @@ "openarm-planner-coordinator": "dimos.robot.manipulators.openarm.blueprints.planner:openarm_planner_coordinator", "openyam-planner-coordinator": "dimos.robot.manipulators.openyam.blueprints.basic:openyam_planner_coordinator", "path-planner-eval": "dimos.navigation.nav_3d.evaluator.blueprints:path_planner_eval", + "teleop-hosted-go2-livekit": "dimos.teleop.hosted.blueprints.livekit:teleop_hosted_go2_livekit", "teleop-hosted-go2-multicam": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_go2_multicam", "teleop-hosted-go2-transport": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_go2_transport", "teleop-hosted-xarm6": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_xarm6", + "teleop-hosted-xarm6-livekit": "dimos.teleop.hosted.blueprints.livekit:teleop_hosted_xarm6_livekit", "teleop-hosted-xarm7": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_xarm7", + "teleop-hosted-xarm7-livekit": "dimos.teleop.hosted.blueprints.livekit:teleop_hosted_xarm7_livekit", "teleop-phone": "dimos.teleop.phone.blueprints:teleop_phone", "teleop-phone-go2": "dimos.teleop.phone.blueprints:teleop_phone_go2", "teleop-phone-go2-fleet": "dimos.teleop.phone.blueprints:teleop_phone_go2_fleet", @@ -224,6 +227,9 @@ "joystick-module": "dimos.robot.unitree.b1.joystick_module.JoystickModule", "keyboard-teleop": "dimos.robot.unitree.keyboard_teleop.KeyboardTeleop", "keyboard-teleop-module": "dimos.teleop.keyboard.keyboard_teleop_module.KeyboardTeleopModule", + "live-kit-front-camera": "dimos.teleop.hosted.blueprints.livekit.LiveKitFrontCamera", + "live-kit-teleop-module": "dimos.teleop.hosted.livekit.LiveKitTeleopModule", + "live-kit-wrist-camera": "dimos.teleop.hosted.blueprints.livekit.LiveKitWristCamera", "local-planner": "dimos.navigation.cmu_nav.modules.local_planner.local_planner.LocalPlanner", "manipulation-module": "dimos.manipulation.manipulation_module.ManipulationModule", "map": "dimos.robot.unitree.type.map.Map", diff --git a/dimos/teleop/hosted/blueprints/livekit.py b/dimos/teleop/hosted/blueprints/livekit.py new file mode 100644 index 0000000000..73a75462fe --- /dev/null +++ b/dimos/teleop/hosted/blueprints/livekit.py @@ -0,0 +1,121 @@ +# 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. + +"""Hosted teleoperation blueprints backed by a LiveKit edge module.""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.transport import LCMTransport +from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.voxels.module import VoxelGridMapper +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.navigation.movement_manager.movement_manager import MovementManager +from dimos.navigation.replanning_a_star.module import ReplanningAStarPlanner +from dimos.robot.manipulators.xarm.blueprints.teleop import ( + coordinator_teleop_xarm6, + coordinator_teleop_xarm7, +) +from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.teleop.hosted.arm_command import ArmCommandModule +from dimos.teleop.hosted.camera_mux import CameraMuxModule +from dimos.teleop.hosted.go2_command import Go2CommandModule +from dimos.teleop.hosted.hosted_stats import HostedStatsModule +from dimos.teleop.hosted.livekit import LiveKitTeleopModule +from dimos.teleop.hosted.map_compress import MapCompressModule +from dimos.teleop.hosted.robot_type import RobotType + + +class LiveKitFrontCamera(RealSenseCamera): + pass + + +class LiveKitWristCamera(RealSenseCamera): + pass + + +teleop_hosted_go2_livekit = ( + autoconnect( + GO2Connection.blueprint(), + Go2CommandModule.blueprint(allow_acrobatics=True), + CameraMuxModule.blueprint(cameras=["cam1"]), + HostedStatsModule.blueprint(), + MapCompressModule.blueprint(), + LiveKitTeleopModule.blueprint(robot_type=RobotType.GO2), + VoxelGridMapper.blueprint(emit_every=5), + CostMapper.blueprint(), + ReplanningAStarPlanner.blueprint(), + MovementManager.blueprint(), + ) + .remappings([(GO2Connection, "color_image", "cam1")]) + .transports( + { + ("tele_cmd_vel", Twist): LCMTransport.spec("/hosted/tele_cmd_vel", Twist), + ("nav_cmd_vel", Twist): LCMTransport.spec("/hosted/nav_cmd_vel", Twist), + ("cmd_vel", Twist): LCMTransport.spec("/hosted/cmd_vel", Twist), + } + ) + .global_config(viewer="none", n_workers=2) +) + + +teleop_hosted_xarm6_livekit = ( + autoconnect( + ArmCommandModule.blueprint(task_names={"right": "teleop_xarm"}), + HostedStatsModule.blueprint(), + CameraMuxModule.blueprint(cameras=["cam1", "cam2"]), + LiveKitTeleopModule.blueprint(robot_type=RobotType.ARM), + coordinator_teleop_xarm6, + LiveKitFrontCamera.blueprint( + camera_name="front", enable_depth=False, enable_pointcloud=False + ), + LiveKitWristCamera.blueprint( + camera_name="wrist", enable_depth=False, enable_pointcloud=False + ), + ) + .remappings( + [ + (LiveKitFrontCamera, "color_image", "cam1"), + (LiveKitWristCamera, "color_image", "cam2"), + (ArmCommandModule, "right_controller_output", "coordinator_cartesian_command"), + ] + ) + .global_config(viewer="none", n_workers=1) +) + + +teleop_hosted_xarm7_livekit = ( + autoconnect( + ArmCommandModule.blueprint(task_names={"right": "teleop_xarm"}), + HostedStatsModule.blueprint(), + CameraMuxModule.blueprint(cameras=["cam1", "cam2"]), + LiveKitTeleopModule.blueprint(robot_type=RobotType.ARM), + coordinator_teleop_xarm7, + LiveKitFrontCamera.blueprint( + camera_name="front", enable_depth=False, enable_pointcloud=False + ), + LiveKitWristCamera.blueprint( + camera_name="wrist", enable_depth=False, enable_pointcloud=False + ), + ) + .remappings( + [ + (LiveKitFrontCamera, "color_image", "cam1"), + (LiveKitWristCamera, "color_image", "cam2"), + (ArmCommandModule, "right_controller_output", "coordinator_cartesian_command"), + ] + ) + .global_config(viewer="none", n_workers=1) +) diff --git a/dimos/teleop/hosted/livekit.py b/dimos/teleop/hosted/livekit.py new file mode 100644 index 0000000000..5084bc4e7e --- /dev/null +++ b/dimos/teleop/hosted/livekit.py @@ -0,0 +1,445 @@ +# 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. + +"""LiveKit edge module for hosted teleoperation. + +One module owns one LiveKit room and translates its data and media planes into +the existing hosted teleoperation ports. It deliberately does not implement a +generic DimOS transport. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from dataclasses import dataclass +import importlib.util +import json +import os +import threading +from typing import TYPE_CHECKING, Any + +import numpy as np +from reactivex.disposable import Disposable + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.teleop.hosted.robot_type import RobotType +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +LIVEKIT_AVAILABLE = ( + importlib.util.find_spec("livekit") is not None + and importlib.util.find_spec("httpx") is not None +) + +if TYPE_CHECKING: + from livekit import rtc + + +@dataclass(frozen=True) +class LiveKitSession: + """Connection material minted by the hosted teleop broker.""" + + session_id: str + url: str + token: str + room: str + + +class LiveKitTeleopConfig(ModuleConfig): + """Broker access and connection policy for one robot LiveKit room.""" + + broker_url: str | None = None + api_key: str | None = None + robot_id: str | None = None + robot_name: str = "robot" + robot_type: RobotType | None = None + operator_identity: str | None = None + heartbeat_hz: float = 1.0 + + +class LiveKitTeleopModule(Module): + """LiveKit room boundary for the hosted teleoperation module graph.""" + + config: LiveKitTeleopConfig + + state_json: Out[bytes] + camera_select: Out[bytes] + cmd_raw: Out[bytes] + cmd_vel_in: Out[TwistStamped] + + mux_image: In[Image] + telemetry_out: In[bytes] + cmd_ack: In[bytes] + map_out: In[bytes] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._loop: asyncio.AbstractEventLoop | None = None + self._worker: threading.Thread | None = None + self._ready = threading.Event() + self._stop_event = threading.Event() + self._lock = threading.RLock() + self._start_error: Exception | None = None + self._room: rtc.Room | None = None + self._video_source: rtc.VideoSource | None = None + self._video_track: rtc.LocalVideoTrack | None = None + self._video_publish_task: asyncio.Task[None] | None = None + self._operator_present = False + self._operator_lost = False + + @rpc + def start(self) -> None: + if not LIVEKIT_AVAILABLE: + raise RuntimeError("livekit and httpx required: pip install dimos[livekit]") + super().start() + self._ready.clear() + self._stop_event.clear() + self._start_error = None + self._register_streams() + self._worker = threading.Thread(target=self._run, daemon=True, name="livekit-teleop") + self._worker.start() + if not self._ready.wait(timeout=30.0): + self.stop() + raise RuntimeError("LiveKit room did not connect within 30 seconds") + if self._start_error is not None: + error = self._start_error + self.stop() + raise RuntimeError("LiveKit room failed to connect") from error + + @rpc + def stop(self) -> None: + self._stop_event.set() + super().stop() + + worker = self._worker + if worker is not None: + worker.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if worker.is_alive(): + logger.error("LiveKit teleop worker did not stop") + else: + self._worker = None + + def _register_streams(self) -> None: + self.register_disposable(Disposable(self.mux_image.subscribe(self._publish_video))) + self.register_disposable( + Disposable( + self.telemetry_out.subscribe( + lambda data: self._publish_data("state_reliable_back", data, reliable=True) + ) + ) + ) + self.register_disposable( + Disposable( + self.cmd_ack.subscribe( + lambda data: self._publish_data("state_reliable_back", data, reliable=True) + ) + ) + ) + self.register_disposable( + Disposable( + self.map_out.subscribe( + lambda data: self._publish_data("map_unreliable", data, reliable=False) + ) + ) + ) + + def _run(self) -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + with self._lock: + self._loop = loop + try: + loop.run_until_complete(self._run_room()) + finally: + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.close() + with self._lock: + self._loop = None + + async def _run_room(self) -> None: + session: LiveKitSession | None = None + heartbeat_task: asyncio.Task[None] | None = None + try: + session = await self._create_session() + await self._connect_room(session) + heartbeat_task = asyncio.create_task(self._heartbeat_loop(session.session_id)) + self._ready.set() + while not self._stop_event.is_set(): + await asyncio.sleep(0.1) + except Exception as error: + self._start_error = error + self._ready.set() + finally: + if heartbeat_task is not None: + heartbeat_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await heartbeat_task + await self._disconnect_room() + if session is not None: + await self._delete_session(session.session_id) + + async def _create_session(self) -> LiveKitSession: + import httpx + + broker_url, api_key = self._broker_credentials() + + payload: dict[str, str] = { + "transport": "livekit", + "robot_name": self.config.robot_name, + } + robot_id = self.config.robot_id or os.environ.get("TELEOP_ROBOT_ID") + if robot_id: + payload["robot_id"] = robot_id + if self.config.robot_type is not None: + payload["robot_type"] = self.config.robot_type.value + + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + f"{broker_url.rstrip('/')}/api/v1/sessions", + headers={"X-Robot-API-Key": api_key, "Content-Type": "application/json"}, + json=payload, + ) + if response.status_code not in (200, 201): + raise RuntimeError( + f"LiveKit broker session create failed: {response.status_code} {response.text[:200]}" + ) + data = response.json() + try: + return LiveKitSession( + session_id=data["session_id"], + url=data["url"], + token=data["token"], + room=data["room"], + ) + except (KeyError, TypeError) as error: + raise RuntimeError( + "LiveKit broker response missing session_id, url, token, or room" + ) from error + + def _broker_credentials(self) -> tuple[str, str]: + broker_url = self.config.broker_url or os.environ.get("TELEOP_BROKER_URL") + api_key = self.config.api_key or os.environ.get("TELEOP_API_KEY") + if not broker_url: + raise RuntimeError("LiveKitTeleopConfig.broker_url or TELEOP_BROKER_URL required") + if not api_key: + raise RuntimeError("LiveKitTeleopConfig.api_key or TELEOP_API_KEY required") + return broker_url.rstrip("/"), api_key + + async def _heartbeat_loop(self, session_id: str) -> None: + import httpx + + broker_url, api_key = self._broker_credentials() + interval = 1.0 / max(self.config.heartbeat_hz, 0.1) + async with httpx.AsyncClient(timeout=30.0) as client: + while True: + try: + response = await client.post( + f"{broker_url}/api/v1/sessions/{session_id}/heartbeat", + headers={"X-Robot-API-Key": api_key, "Content-Type": "application/json"}, + json={}, + ) + if response.status_code != 200: + logger.warning("LiveKit heartbeat failed", status=response.status_code) + except Exception: + logger.warning("LiveKit heartbeat failed", exc_info=True) + await asyncio.sleep(interval) + + async def _delete_session(self, session_id: str) -> None: + import httpx + + broker_url, api_key = self._broker_credentials() + try: + async with httpx.AsyncClient(timeout=30.0) as client: + await client.delete( + f"{broker_url}/api/v1/sessions/{session_id}", + headers={"X-Robot-API-Key": api_key, "Content-Type": "application/json"}, + ) + except Exception: + logger.warning("LiveKit broker session delete failed", exc_info=True) + + async def _connect_room(self, session: LiveKitSession) -> None: + from livekit import rtc + + room = rtc.Room() + with self._lock: + self._room = room + + @room.on("data_received") # type: ignore[untyped-decorator] + def _data_received(packet: Any) -> None: + self._on_data(getattr(packet, "topic", "") or "", bytes(packet.data)) + + @room.on("participant_connected") # type: ignore[untyped-decorator] + def _participant_connected(participant: Any) -> None: + if self._is_operator(participant): + with self._lock: + self._operator_present = True + self._operator_lost = False + + @room.on("participant_disconnected") # type: ignore[untyped-decorator] + def _participant_disconnected(participant: Any) -> None: + identity = getattr(participant, "identity", None) + if self._is_operator(participant) and not self._has_operator( + room, excluded_identity=identity + ): + self._on_operator_lost() + + @room.on("disconnected") # type: ignore[untyped-decorator] + def _disconnected(_reason: Any) -> None: + self._on_operator_lost() + + await room.connect(session.url, session.token) + with self._lock: + self._operator_present = self._has_operator(room) + self._operator_lost = False + logger.info("LiveKit teleop connected", room=session.room) + + async def _disconnect_room(self) -> None: + with self._lock: + room = self._room + self._room = None + self._video_source = None + self._video_track = None + self._video_publish_task = None + if room is not None: + with contextlib.suppress(Exception): + await room.disconnect() + + def _is_operator(self, participant: Any) -> bool: + identity = getattr(participant, "identity", None) + return self.config.operator_identity is None or identity == self.config.operator_identity + + def _has_operator(self, room: rtc.Room, excluded_identity: str | None = None) -> bool: + return any( + self._is_operator(participant) + and getattr(participant, "identity", None) != excluded_identity + for participant in room.remote_participants.values() + ) + + def _publish_data(self, topic: str, data: bytes, reliable: bool) -> None: + with self._lock: + loop, room = self._loop, self._room + if loop is None or room is None or not loop.is_running(): + return + future = asyncio.run_coroutine_threadsafe( + room.local_participant.publish_data(bytes(data), reliable=reliable, topic=topic), loop + ) + future.add_done_callback(self._log_publish_failure) + + def _publish_video(self, image: Image) -> None: + with self._lock: + loop = self._loop + if loop is None or not loop.is_running(): + return + try: + width, height, rgba = self._image_to_rgba(image) + except Exception: + logger.warning("LiveKit video frame conversion failed", exc_info=True) + return + loop.call_soon_threadsafe(self._capture_video, width, height, rgba) + + def _on_data(self, topic: str, data: bytes) -> None: + if topic == "state_reliable": + self.state_json.publish(data) + self.camera_select.publish(data) + elif topic == "cmd_unreliable": + self.cmd_raw.publish(data) + with contextlib.suppress(Exception): + self.cmd_vel_in.publish(TwistStamped.lcm_decode(data)) + + def _on_operator_lost(self) -> None: + with self._lock: + if self._stop_event.is_set() or not self._operator_present or self._operator_lost: + return + self._operator_lost = True + self._operator_present = False + logger.warning("LiveKit operator link lost") + self.state_json.publish(json.dumps({"type": "operator_lost"}).encode()) + + def _capture_video(self, width: int, height: int, rgba: bytes) -> None: + from livekit import rtc + + if self._video_source is None: + self._video_source = rtc.VideoSource(width, height) + self._video_publish_task = asyncio.create_task(self._publish_video_track()) + self._video_source.capture_frame( + rtc.VideoFrame(width, height, rtc.VideoBufferType.RGBA, rgba) + ) + + async def _publish_video_track(self) -> None: + from livekit import rtc + + with self._lock: + room, source = self._room, self._video_source + if room is None or source is None: + return + try: + track = rtc.LocalVideoTrack.create_video_track("camera", source) + await room.local_participant.publish_track( + track, rtc.TrackPublishOptions(source=rtc.TrackSource.SOURCE_CAMERA) + ) + except Exception: + with self._lock: + self._video_source = None + self._video_track = None + self._video_publish_task = None + logger.warning("LiveKit video track publish failed", exc_info=True) + return + with self._lock: + self._video_track = track + + @staticmethod + def _image_to_rgba(image: Image) -> tuple[int, int, bytes]: + data = image.data + if data.dtype == np.uint16: + data = (data >> 8).astype(np.uint8) + elif data.dtype != np.uint8: + data = data.astype(np.uint8) + height, width = data.shape[:2] + if image.format == ImageFormat.RGBA: + rgba = data + elif image.format == ImageFormat.BGRA: + rgba = data[..., [2, 1, 0, 3]] + elif image.format == ImageFormat.RGB: + rgba = np.dstack([data, np.full((height, width), 255, dtype=np.uint8)]) + elif image.format in (ImageFormat.GRAY, ImageFormat.GRAY16): + gray = data if data.ndim == 2 else data[..., 0] + rgba = np.dstack([gray, gray, gray, np.full((height, width), 255, dtype=np.uint8)]) + else: + rgba = np.dstack( + [ + data[..., 2], + data[..., 1], + data[..., 0], + np.full((height, width), 255, dtype=np.uint8), + ] + ) + return width, height, np.ascontiguousarray(rgba).tobytes() + + @staticmethod + def _log_publish_failure(future: Any) -> None: + with contextlib.suppress(asyncio.CancelledError): + try: + future.result() + except Exception: + logger.warning("LiveKit data publish failed", exc_info=True) diff --git a/dimos/teleop/hosted/test_livekit.py b/dimos/teleop/hosted/test_livekit.py new file mode 100644 index 0000000000..3042088db7 --- /dev/null +++ b/dimos/teleop/hosted/test_livekit.py @@ -0,0 +1,215 @@ +# 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 __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import numpy as np +import pytest + +from dimos.core.module import Module +from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.teleop.hosted.blueprints.livekit import ( + teleop_hosted_go2_livekit, + teleop_hosted_xarm6_livekit, + teleop_hosted_xarm7_livekit, +) +from dimos.teleop.hosted.livekit import LiveKitSession, LiveKitTeleopConfig, LiveKitTeleopModule +from dimos.teleop.hosted.robot_type import RobotType + + +class LiveKitTestModule(LiveKitTeleopModule): + state_json: MagicMock + camera_select: MagicMock + cmd_raw: MagicMock + cmd_vel_in: MagicMock + mux_image: MagicMock + telemetry_out: MagicMock + cmd_ack: MagicMock + map_out: MagicMock + + +@pytest.fixture +def module(monkeypatch: pytest.MonkeyPatch) -> LiveKitTestModule: + monkeypatch.setattr( + Module, + "__init__", + lambda self, **kwargs: setattr(self, "config", LiveKitTeleopConfig(**kwargs)), + ) + result = LiveKitTestModule( + broker_url="https://broker.example", + api_key="robot-key", + robot_type=RobotType.GO2, + ) + for name in ( + "state_json", + "camera_select", + "cmd_raw", + "cmd_vel_in", + "mux_image", + "telemetry_out", + "cmd_ack", + "map_out", + ): + setattr(result, name, MagicMock()) + return result + + +def test_state_packets_fan_out_to_command_and_camera_selection(module: LiveKitTestModule) -> None: + module._on_data("state_reliable", b'{"type":"camera_select"}') + + module.state_json.publish.assert_called_once_with(b'{"type":"camera_select"}') + module.camera_select.publish.assert_called_once_with(b'{"type":"camera_select"}') + + +def test_command_packets_reach_raw_and_typed_command_ports( + module: LiveKitTestModule, mocker: pytest.MockFixture +) -> None: + command = MagicMock(spec=TwistStamped) + mocker.patch.object(TwistStamped, "lcm_decode", return_value=command) + + module._on_data("cmd_unreliable", b"command") + + module.cmd_raw.publish.assert_called_once_with(b"command") + module.cmd_vel_in.publish.assert_called_once_with(command) + + +def test_non_twist_command_still_reaches_raw_port( + module: LiveKitTestModule, mocker: pytest.MockFixture +) -> None: + mocker.patch.object(TwistStamped, "lcm_decode", side_effect=ValueError("not a twist")) + + module._on_data("cmd_unreliable", b"arm-command") + + module.cmd_raw.publish.assert_called_once_with(b"arm-command") + module.cmd_vel_in.publish.assert_not_called() + + +def test_operator_loss_is_emitted_once(module: LiveKitTestModule) -> None: + module._operator_present = True + + module._on_operator_lost() + module._on_operator_lost() + + module.state_json.publish.assert_called_once_with(b'{"type": "operator_lost"}') + + +def test_broker_session_request_uses_robot_metadata( + module: LiveKitTestModule, mocker: pytest.MockFixture +) -> None: + response = MagicMock(status_code=201) + response.json.return_value = { + "session_id": "session", + "url": "wss://livekit.example", + "token": "jwt", + "room": "room", + } + client = MagicMock() + client.post = AsyncMock(return_value=response) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + mocker.patch("httpx.AsyncClient", return_value=client) + + session = asyncio.run(module._create_session()) + + assert session == LiveKitSession( + session_id="session", url="wss://livekit.example", token="jwt", room="room" + ) + assert client.post.await_args.kwargs["json"] == { + "transport": "livekit", + "robot_name": "robot", + "robot_type": "go2", + } + + +def test_outbound_data_preserves_topic_reliability( + module: LiveKitTestModule, mocker: pytest.MockFixture +) -> None: + room = MagicMock() + room.local_participant.publish_data = AsyncMock() + loop = MagicMock() + loop.is_running.return_value = True + future = MagicMock() + submit = mocker.patch("asyncio.run_coroutine_threadsafe", return_value=future) + module._room = room + module._loop = loop + + module._publish_data("map_unreliable", b"map", reliable=False) + + coro = submit.call_args.args[0] + assert submit.call_args.args[1] is loop + coro.close() + room.local_participant.publish_data.assert_called_once_with( + b"map", reliable=False, topic="map_unreliable" + ) + future.add_done_callback.assert_called_once() + + +def test_image_conversion_produces_rgba() -> None: + image = MagicMock(spec=Image) + image.data = np.array([[[1, 2, 3]]], dtype=np.uint8) + image.format = ImageFormat.BGR + + width, height, rgba = LiveKitTeleopModule._image_to_rgba(image) + + assert (width, height) == (1, 1) + assert rgba == bytes([3, 2, 1, 255]) + + +def test_livekit_blueprints_are_exposed() -> None: + assert teleop_hosted_go2_livekit is not None + assert teleop_hosted_xarm6_livekit is not None + assert teleop_hosted_xarm7_livekit is not None + + +def test_connect_registers_data_and_operator_lifecycle( + module: LiveKitTestModule, mocker: pytest.MockFixture +) -> None: + room = MagicMock() + room.connect = AsyncMock() + room.remote_participants = {} + handlers: dict[str, object] = {} + + def on(event: str): + def register(handler: object) -> object: + handlers[event] = handler + return handler + + return register + + room.on.side_effect = on + mocker.patch("livekit.rtc.Room", return_value=room) + + asyncio.run( + module._connect_room(LiveKitSession("session", "wss://livekit.example", "jwt", "room")) + ) + + assert set(handlers) == { + "data_received", + "participant_connected", + "participant_disconnected", + "disconnected", + } + handlers["data_received"](SimpleNamespace(topic="state_reliable", data=b"state")) # type: ignore[operator] + room.remote_participants = {"operator": SimpleNamespace(identity="operator")} + module._operator_present = True + handlers["participant_disconnected"](SimpleNamespace(identity="operator")) # type: ignore[operator] + assert [call.args[0] for call in module.state_json.publish.call_args_list] == [ + b"state", + b'{"type": "operator_lost"}', + ] diff --git a/pyproject.toml b/pyproject.toml index 1fc0193bab..b686017a19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -342,6 +342,12 @@ webrtc = [ "aiohttp>=3.9.0", ] +livekit = [ + # Hosted teleoperation room client. + "livekit>=1.0.0", + "httpx>=0.27.0", +] + base = [ "dimos[agents,web,perception,visualization]", ] @@ -374,7 +380,7 @@ graspgenx = [ ] all = [ - "dimos[agents,apriltag,base,cpu,cuda,drone,manipulation,misc,perception,scene,sim,unitree,visualization,web,webrtc]", + "dimos[agents,apriltag,base,cpu,cuda,drone,livekit,manipulation,misc,perception,scene,sim,unitree,visualization,web,webrtc]", ] [dependency-groups] @@ -383,7 +389,7 @@ autofix = ["ruff==0.14.3"] # Project deps shared by `tests` and `lint`. project-deps = [ - "dimos[web,visualization,webrtc]", + "dimos[livekit,web,visualization,webrtc]", "torch", "langchain==1.2.3", "langchain-core==1.3.3", @@ -608,6 +614,8 @@ module = [ "h5py.*", "mcap", "mcap.*", + "livekit", + "livekit.*", "mujoco", "mujoco_playground.*", "nav_msgs.*", diff --git a/uv.lock b/uv.lock index e2313b7058..41e7025b51 100644 --- a/uv.lock +++ b/uv.lock @@ -26,7 +26,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-07-25T20:06:13.81106606Z" exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -1661,6 +1661,7 @@ all = [ { name = "gdown" }, { name = "googlemaps" }, { name = "gtsam-extended" }, + { name = "httpx" }, { name = "hydra-core" }, { name = "ipykernel" }, { name = "jinja2" }, @@ -1670,6 +1671,7 @@ all = [ { name = "langchain-ollama" }, { name = "langchain-openai" }, { name = "lap" }, + { name = "livekit" }, { name = "manifold3d" }, { name = "matplotlib" }, { name = "moondream" }, @@ -1771,6 +1773,10 @@ learning = [ { name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pyarrow" }, ] +livekit = [ + { name = "httpx" }, + { name = "livekit" }, +] manipulation = [ { name = "a750-control", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "drake", version = "1.45.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'aarch64' and sys_platform == 'darwin'" }, @@ -1920,7 +1926,7 @@ browser-tests = [ lint = [ { name = "aiortc" }, { name = "chromadb" }, - { name = "dimos", extra = ["visualization", "web", "webrtc"] }, + { name = "dimos", extra = ["livekit", "visualization", "web", "webrtc"] }, { name = "einops" }, { name = "gdown" }, { name = "googlemaps" }, @@ -1959,7 +1965,7 @@ lint = [ ] project-deps = [ { name = "chromadb" }, - { name = "dimos", extra = ["visualization", "web", "webrtc"] }, + { name = "dimos", extra = ["livekit", "visualization", "web", "webrtc"] }, { name = "einops" }, { name = "gdown" }, { name = "googlemaps" }, @@ -1983,7 +1989,7 @@ tests = [ { name = "chromadb" }, { name = "coacd" }, { name = "coverage" }, - { name = "dimos", extra = ["apriltag", "cpu", "drone", "learning", "mapping", "visualization", "web", "webrtc"] }, + { name = "dimos", extra = ["apriltag", "cpu", "drone", "learning", "livekit", "mapping", "visualization", "web", "webrtc"] }, { name = "einops" }, { name = "gdown" }, { name = "googlemaps" }, @@ -2030,7 +2036,7 @@ tests-self-hosted = [ { name = "chromadb" }, { name = "coacd" }, { name = "coverage" }, - { name = "dimos", extra = ["agents", "apriltag", "cpu", "drone", "learning", "manipulation", "mapping", "misc", "perception", "sim", "unitree", "visualization", "web", "webrtc"] }, + { name = "dimos", extra = ["agents", "apriltag", "cpu", "drone", "learning", "livekit", "manipulation", "mapping", "misc", "perception", "sim", "unitree", "visualization", "web", "webrtc"] }, { name = "einops" }, { name = "gdown" }, { name = "googlemaps" }, @@ -2092,7 +2098,7 @@ requires-dist = [ { name = "cupy-cuda12x", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = "==13.6.0" }, { name = "cyclonedds", marker = "extra == 'dds'", specifier = ">=0.10.5" }, { name = "cyclonedds", marker = "extra == 'unitree-dds'", specifier = ">=0.10.5" }, - { name = "dimos", extras = ["agents", "apriltag", "base", "cpu", "cuda", "drone", "manipulation", "misc", "perception", "scene", "sim", "unitree", "visualization", "web", "webrtc"], marker = "extra == 'all'" }, + { name = "dimos", extras = ["agents", "apriltag", "base", "cpu", "cuda", "drone", "livekit", "manipulation", "misc", "perception", "scene", "sim", "unitree", "visualization", "web", "webrtc"], marker = "extra == 'all'" }, { name = "dimos", extras = ["agents", "web", "perception", "visualization"], marker = "extra == 'base'" }, { name = "dimos", extras = ["base", "mapping"], marker = "extra == 'unitree'" }, { name = "dimos", extras = ["unitree"], marker = "extra == 'unitree-dds'" }, @@ -2114,6 +2120,7 @@ requires-dist = [ { name = "gtsam-extended", marker = "extra == 'mapping'", specifier = ">=4.3a1.post1" }, { name = "h5py", marker = "extra == 'learning'" }, { name = "huggingface-hub", marker = "extra == 'graspgenx'", specifier = ">=0.30,<1" }, + { name = "httpx", marker = "extra == 'livekit'", specifier = ">=0.27.0" }, { name = "hydra-core", marker = "extra == 'perception'", specifier = ">=1.3.0" }, { name = "imagecodecs", specifier = ">=2024.6.1" }, { name = "ipykernel", marker = "extra == 'misc'" }, @@ -2126,6 +2133,7 @@ requires-dist = [ { name = "langchain-openai", marker = "extra == 'agents'", specifier = ">=1,<2" }, { name = "lap", marker = "extra == 'perception'", specifier = ">=0.5.12" }, { name = "lazy-loader" }, + { name = "livekit", marker = "extra == 'livekit'", specifier = ">=1.0.0" }, { name = "llvmlite", specifier = ">=0.42.0" }, { name = "lz4", specifier = ">=4.4.5" }, { name = "manifold3d", marker = "extra == 'apriltag'", specifier = ">=2.5.0" }, @@ -2209,7 +2217,7 @@ requires-dist = [ { name = "yourdfpy", marker = "(platform_machine != 'aarch64' and extra == 'visualization') or (sys_platform != 'linux' and extra == 'visualization')", specifier = ">=0.0.60" }, { name = "yourdfpy", marker = "extra == 'manipulation'", specifier = ">=0.0.60" }, ] -provides-extras = ["misc", "visualization", "learning", "agents", "web", "perception", "unitree", "unitree-dds", "manipulation", "cpu", "cuda", "sim", "mapping", "drone", "dds", "webrtc", "base", "apriltag", "scene", "graspgenx", "all"] +provides-extras = ["misc", "visualization", "learning", "agents", "web", "perception", "unitree", "unitree-dds", "manipulation", "cpu", "cuda", "sim", "mapping", "drone", "dds", "webrtc", "livekit", "base", "apriltag", "scene", "graspgenx", "all"] [package.metadata.requires-dev] autofix = [{ name = "ruff", specifier = "==0.14.3" }] @@ -2217,7 +2225,7 @@ browser-tests = [{ name = "playwright", specifier = ">=1.55" }] lint = [ { name = "aiortc", specifier = ">=1.14.0" }, { name = "chromadb", specifier = ">=1.0.0" }, - { name = "dimos", extras = ["web", "visualization", "webrtc"] }, + { name = "dimos", extras = ["livekit", "web", "visualization", "webrtc"] }, { name = "einops", specifier = ">=0.8.1" }, { name = "gdown", specifier = "==6.0.0" }, { name = "googlemaps", specifier = ">=4.10.0" }, @@ -2255,7 +2263,7 @@ lint = [ ] project-deps = [ { name = "chromadb", specifier = ">=1.0.0" }, - { name = "dimos", extras = ["web", "visualization", "webrtc"] }, + { name = "dimos", extras = ["livekit", "web", "visualization", "webrtc"] }, { name = "einops", specifier = ">=0.8.1" }, { name = "gdown", specifier = "==6.0.0" }, { name = "googlemaps", specifier = ">=4.10.0" }, @@ -2280,7 +2288,7 @@ tests = [ { name = "coacd", specifier = ">=1.0.0" }, { name = "coverage", specifier = ">=7.0" }, { name = "dimos", extras = ["apriltag", "mapping", "drone", "cpu", "learning"] }, - { name = "dimos", extras = ["web", "visualization", "webrtc"] }, + { name = "dimos", extras = ["livekit", "web", "visualization", "webrtc"] }, { name = "einops", specifier = ">=0.8.1" }, { name = "gdown", specifier = "==6.0.0" }, { name = "googlemaps", specifier = ">=4.10.0" }, @@ -2329,7 +2337,7 @@ tests-self-hosted = [ { name = "coverage", specifier = ">=7.0" }, { name = "dimos", extras = ["agents", "perception", "manipulation", "sim", "unitree", "misc"] }, { name = "dimos", extras = ["apriltag", "mapping", "drone", "cpu", "learning"] }, - { name = "dimos", extras = ["web", "visualization", "webrtc"] }, + { name = "dimos", extras = ["livekit", "web", "visualization", "webrtc"] }, { name = "einops", specifier = ">=0.8.1" }, { name = "gdown", specifier = "==6.0.0" }, { name = "googlemaps", specifier = ">=4.10.0" }, @@ -4560,6 +4568,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/1e/b832de447dee8b582cac175871d2f6c3d5077cc56d5575cadba1fd1cccfa/linkify_it_py-2.0.3-py3-none-any.whl", hash = "sha256:6bcbc417b0ac14323382aef5c5192c0075bf8a9d6b41820a2b66371eac6b6d79", size = 19820, upload-time = "2024-02-04T14:48:02.496Z" }, ] +[[package]] +name = "livekit" +version = "1.1.13" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiofiles" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "protobuf" }, + { name = "types-protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8d/92/dcd4f295913533ddd0d48153cc28f1358d550bea651460bd895256981c4d/livekit-1.1.13.tar.gz", hash = "sha256:aa2bd89cf0c2ebcaa71a240275964c23900a0422a2a3d43d274e88a211a0ecfc", size = 370211, upload-time = "2026-06-30T11:54:00.321Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/bd/15100217109595aedbb9bcfdfc1c77513c0f44940d72644d16b611476941/livekit-1.1.13-py3-none-macosx_10_15_x86_64.whl", hash = "sha256:2a19b023de9a573fe5da629e3ee514f2962c87e1afec95a6b19bbbdb6ea8f703", size = 10147642, upload-time = "2026-06-30T11:53:48.781Z" }, + { url = "https://files.pythonhosted.org/packages/97/dd/4f001a9c5ccde361a53437a09bc04dc9cad8003c6711c2bcc0734e18e626/livekit-1.1.13-py3-none-macosx_11_0_arm64.whl", hash = "sha256:90b6796e0c4515bc1e8a1e109a40b88d82c36b12452b7ae01fe35be7ee8357ba", size = 8968740, upload-time = "2026-06-30T11:53:51.223Z" }, + { url = "https://files.pythonhosted.org/packages/26/1a/7e97a45a4b6e10ce3a4e9938e3d3fa0a6512a6557f5990685eab4f6e88d1/livekit-1.1.13-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:43779c0f3bb27589cd517d60c442c674c2f967a97db158829a533974a29a3997", size = 9980507, upload-time = "2026-06-30T11:53:53.349Z" }, + { url = "https://files.pythonhosted.org/packages/e3/9d/389bbdf39ccd2c464a4749ba7be1584dea608518c32a5ddbc511db6b0cf7/livekit-1.1.13-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:f4e83f0e272f4b2e9cc39dbefd7bb23b75bd8c105478d867c2869254ca4e149a", size = 11367692, upload-time = "2026-06-30T11:53:55.409Z" }, + { url = "https://files.pythonhosted.org/packages/da/ce/a3d3e0566dbd2586c325240d44afec6d44421eb794bcf8dbaca15463a7b7/livekit-1.1.13-py3-none-win_amd64.whl", hash = "sha256:22dff7a39cb3d590a4757e20d3ce5d326ab882350386f5070f45cbd6d9ccf839", size = 10717013, upload-time = "2026-06-30T11:53:57.701Z" }, +] + [[package]] name = "llvmlite" version = "0.46.0" @@ -9385,6 +9413,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d5/91/9b286ab899c008c2cb05e8be99814807e7fbbd33f0c0c960470826e5ac82/typer-0.23.1-py3-none-any.whl", hash = "sha256:3291ad0d3c701cbf522012faccfbb29352ff16ad262db2139e6b01f15781f14e", size = 56813, upload-time = "2026-02-13T10:04:32.008Z" }, ] +[[package]] +name = "types-protobuf" +version = "7.34.1.20260518" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/29/59/e2b13b499d15e6720150c4b1a8d91e31fcacf716b432397475b3151ff7e4/types_protobuf-7.34.1.20260518.tar.gz", hash = "sha256:28cfaded25889cb83ebfb63cfb0a43628f0b6f3785767bec17287dc6468795f2", size = 68936, upload-time = "2026-05-18T06:01:47.332Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/1f/ec5caf72c2e3b688ca3927e0979a04ddad19e1afc4bf1c199bd743e0f419/types_protobuf-7.34.1.20260518-py3-none-any.whl", hash = "sha256:a0a5337413347166439c0e07cbc26c6164d091401c6f01b1dfd8cdb966c4dd8f", size = 85992, upload-time = "2026-05-18T06:01:45.696Z" }, +] + [[package]] name = "types-pyaudio" version = "0.2.16.20260508" From 051ae808bf68004caafb414b3c83ac0317bbc901 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 7 Aug 2026 20:01:00 -0700 Subject: [PATCH 2/4] refactor(teleop): extract LiveKit broker client --- dimos/teleop/hosted/livekit.py | 103 ++--------------- dimos/teleop/hosted/livekit_broker_client.py | 112 +++++++++++++++++++ dimos/teleop/hosted/test_livekit.py | 9 +- 3 files changed, 128 insertions(+), 96 deletions(-) create mode 100644 dimos/teleop/hosted/livekit_broker_client.py diff --git a/dimos/teleop/hosted/livekit.py b/dimos/teleop/hosted/livekit.py index 5084bc4e7e..e6fffe0ae9 100644 --- a/dimos/teleop/hosted/livekit.py +++ b/dimos/teleop/hosted/livekit.py @@ -23,10 +23,8 @@ import asyncio import contextlib -from dataclasses import dataclass import importlib.util import json -import os import threading from typing import TYPE_CHECKING, Any @@ -39,6 +37,7 @@ from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.TwistStamped import TwistStamped from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.teleop.hosted.livekit_broker_client import LiveKitBrokerClient, LiveKitSession from dimos.teleop.hosted.robot_type import RobotType from dimos.utils.logging_config import setup_logger @@ -53,16 +52,6 @@ from livekit import rtc -@dataclass(frozen=True) -class LiveKitSession: - """Connection material minted by the hosted teleop broker.""" - - session_id: str - url: str - token: str - room: str - - class LiveKitTeleopConfig(ModuleConfig): """Broker access and connection policy for one robot LiveKit room.""" @@ -104,6 +93,7 @@ def __init__(self, **kwargs: Any) -> None: self._video_publish_task: asyncio.Task[None] | None = None self._operator_present = False self._operator_lost = False + self._broker = LiveKitBrokerClient(self.config.broker_url, self.config.api_key) @rpc def start(self) -> None: @@ -181,9 +171,13 @@ async def _run_room(self) -> None: session: LiveKitSession | None = None heartbeat_task: asyncio.Task[None] | None = None try: - session = await self._create_session() + session = await self._broker.create_session( + self.config.robot_id, self.config.robot_name, self.config.robot_type + ) await self._connect_room(session) - heartbeat_task = asyncio.create_task(self._heartbeat_loop(session.session_id)) + heartbeat_task = asyncio.create_task( + self._broker.heartbeat(session.session_id, self.config.heartbeat_hz) + ) self._ready.set() while not self._stop_event.is_set(): await asyncio.sleep(0.1) @@ -197,86 +191,7 @@ async def _run_room(self) -> None: await heartbeat_task await self._disconnect_room() if session is not None: - await self._delete_session(session.session_id) - - async def _create_session(self) -> LiveKitSession: - import httpx - - broker_url, api_key = self._broker_credentials() - - payload: dict[str, str] = { - "transport": "livekit", - "robot_name": self.config.robot_name, - } - robot_id = self.config.robot_id or os.environ.get("TELEOP_ROBOT_ID") - if robot_id: - payload["robot_id"] = robot_id - if self.config.robot_type is not None: - payload["robot_type"] = self.config.robot_type.value - - async with httpx.AsyncClient(timeout=30.0) as client: - response = await client.post( - f"{broker_url.rstrip('/')}/api/v1/sessions", - headers={"X-Robot-API-Key": api_key, "Content-Type": "application/json"}, - json=payload, - ) - if response.status_code not in (200, 201): - raise RuntimeError( - f"LiveKit broker session create failed: {response.status_code} {response.text[:200]}" - ) - data = response.json() - try: - return LiveKitSession( - session_id=data["session_id"], - url=data["url"], - token=data["token"], - room=data["room"], - ) - except (KeyError, TypeError) as error: - raise RuntimeError( - "LiveKit broker response missing session_id, url, token, or room" - ) from error - - def _broker_credentials(self) -> tuple[str, str]: - broker_url = self.config.broker_url or os.environ.get("TELEOP_BROKER_URL") - api_key = self.config.api_key or os.environ.get("TELEOP_API_KEY") - if not broker_url: - raise RuntimeError("LiveKitTeleopConfig.broker_url or TELEOP_BROKER_URL required") - if not api_key: - raise RuntimeError("LiveKitTeleopConfig.api_key or TELEOP_API_KEY required") - return broker_url.rstrip("/"), api_key - - async def _heartbeat_loop(self, session_id: str) -> None: - import httpx - - broker_url, api_key = self._broker_credentials() - interval = 1.0 / max(self.config.heartbeat_hz, 0.1) - async with httpx.AsyncClient(timeout=30.0) as client: - while True: - try: - response = await client.post( - f"{broker_url}/api/v1/sessions/{session_id}/heartbeat", - headers={"X-Robot-API-Key": api_key, "Content-Type": "application/json"}, - json={}, - ) - if response.status_code != 200: - logger.warning("LiveKit heartbeat failed", status=response.status_code) - except Exception: - logger.warning("LiveKit heartbeat failed", exc_info=True) - await asyncio.sleep(interval) - - async def _delete_session(self, session_id: str) -> None: - import httpx - - broker_url, api_key = self._broker_credentials() - try: - async with httpx.AsyncClient(timeout=30.0) as client: - await client.delete( - f"{broker_url}/api/v1/sessions/{session_id}", - headers={"X-Robot-API-Key": api_key, "Content-Type": "application/json"}, - ) - except Exception: - logger.warning("LiveKit broker session delete failed", exc_info=True) + await self._broker.close_session(session.session_id) async def _connect_room(self, session: LiveKitSession) -> None: from livekit import rtc diff --git a/dimos/teleop/hosted/livekit_broker_client.py b/dimos/teleop/hosted/livekit_broker_client.py new file mode 100644 index 0000000000..3b50ed454b --- /dev/null +++ b/dimos/teleop/hosted/livekit_broker_client.py @@ -0,0 +1,112 @@ +# 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. + +"""Client for broker-minted LiveKit robot sessions.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +import os + +from dimos.teleop.hosted.robot_type import RobotType +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +@dataclass(frozen=True) +class LiveKitSession: + """Connection material minted by the hosted teleop broker.""" + + session_id: str + url: str + token: str + room: str + + +class LiveKitBrokerClient: + """Create, maintain, and close one broker-mediated LiveKit session.""" + + def __init__(self, broker_url: str | None, api_key: str | None) -> None: + self._broker_url = (broker_url or os.environ.get("TELEOP_BROKER_URL", "")).rstrip("/") + self._api_key = api_key or os.environ.get("TELEOP_API_KEY", "") + if not self._broker_url: + raise RuntimeError("broker_url or TELEOP_BROKER_URL required") + if not self._api_key: + raise RuntimeError("api_key or TELEOP_API_KEY required") + + @property + def _headers(self) -> dict[str, str]: + return {"X-Robot-API-Key": self._api_key, "Content-Type": "application/json"} + + async def create_session( + self, robot_id: str | None, robot_name: str, robot_type: RobotType | None + ) -> LiveKitSession: + import httpx + + payload: dict[str, str] = {"transport": "livekit", "robot_name": robot_name} + if robot_id: + payload["robot_id"] = robot_id + if robot_type is not None: + payload["robot_type"] = robot_type.value + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post( + f"{self._broker_url}/api/v1/sessions", headers=self._headers, json=payload + ) + if response.status_code not in (200, 201): + raise RuntimeError( + f"LiveKit broker session create failed: {response.status_code} {response.text[:200]}" + ) + data = response.json() + try: + return LiveKitSession( + session_id=data["session_id"], + url=data["url"], + token=data["token"], + room=data["room"], + ) + except (KeyError, TypeError) as error: + raise RuntimeError( + "LiveKit broker response missing session_id, url, token, or room" + ) from error + + async def heartbeat(self, session_id: str, heartbeat_hz: float) -> None: + import httpx + + interval = 1.0 / max(heartbeat_hz, 0.1) + async with httpx.AsyncClient(timeout=30.0) as client: + while True: + try: + response = await client.post( + f"{self._broker_url}/api/v1/sessions/{session_id}/heartbeat", + headers=self._headers, + json={}, + ) + if response.status_code != 200: + logger.warning("LiveKit heartbeat failed", status=response.status_code) + except Exception: + logger.warning("LiveKit heartbeat failed", exc_info=True) + await asyncio.sleep(interval) + + async def close_session(self, session_id: str) -> None: + import httpx + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + await client.delete( + f"{self._broker_url}/api/v1/sessions/{session_id}", headers=self._headers + ) + except Exception: + logger.warning("LiveKit broker session delete failed", exc_info=True) diff --git a/dimos/teleop/hosted/test_livekit.py b/dimos/teleop/hosted/test_livekit.py index 3042088db7..5940d9873c 100644 --- a/dimos/teleop/hosted/test_livekit.py +++ b/dimos/teleop/hosted/test_livekit.py @@ -29,7 +29,8 @@ teleop_hosted_xarm6_livekit, teleop_hosted_xarm7_livekit, ) -from dimos.teleop.hosted.livekit import LiveKitSession, LiveKitTeleopConfig, LiveKitTeleopModule +from dimos.teleop.hosted.livekit import LiveKitTeleopConfig, LiveKitTeleopModule +from dimos.teleop.hosted.livekit_broker_client import LiveKitSession from dimos.teleop.hosted.robot_type import RobotType @@ -125,7 +126,11 @@ def test_broker_session_request_uses_robot_metadata( client.__aexit__ = AsyncMock(return_value=None) mocker.patch("httpx.AsyncClient", return_value=client) - session = asyncio.run(module._create_session()) + session = asyncio.run( + module._broker.create_session( + module.config.robot_id, module.config.robot_name, module.config.robot_type + ) + ) assert session == LiveKitSession( session_id="session", url="wss://livekit.example", token="jwt", room="room" From f5e3c5c579ac4c1e7b9e5aba568114280cd29c5c Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 7 Aug 2026 20:04:29 -0700 Subject: [PATCH 3/4] fix(teleop): defer LiveKit video until connected --- dimos/teleop/hosted/livekit.py | 6 ++++-- dimos/teleop/hosted/test_livekit.py | 9 +++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/dimos/teleop/hosted/livekit.py b/dimos/teleop/hosted/livekit.py index e6fffe0ae9..20e4d2ee67 100644 --- a/dimos/teleop/hosted/livekit.py +++ b/dimos/teleop/hosted/livekit.py @@ -263,8 +263,8 @@ def _publish_data(self, topic: str, data: bytes, reliable: bool) -> None: def _publish_video(self, image: Image) -> None: with self._lock: - loop = self._loop - if loop is None or not loop.is_running(): + loop, room = self._loop, self._room + if loop is None or room is None or not loop.is_running(): return try: width, height, rgba = self._image_to_rgba(image) @@ -294,6 +294,8 @@ def _on_operator_lost(self) -> None: def _capture_video(self, width: int, height: int, rgba: bytes) -> None: from livekit import rtc + if self._room is None: + return if self._video_source is None: self._video_source = rtc.VideoSource(width, height) self._video_publish_task = asyncio.create_task(self._publish_video_track()) diff --git a/dimos/teleop/hosted/test_livekit.py b/dimos/teleop/hosted/test_livekit.py index 5940d9873c..759efc5902 100644 --- a/dimos/teleop/hosted/test_livekit.py +++ b/dimos/teleop/hosted/test_livekit.py @@ -176,6 +176,15 @@ def test_image_conversion_produces_rgba() -> None: assert rgba == bytes([3, 2, 1, 255]) +def test_video_frames_are_dropped_until_the_room_connects(module: LiveKitTestModule) -> None: + module._loop = MagicMock() + module._loop.is_running.return_value = True + + module._publish_video(MagicMock(spec=Image)) + + assert module._loop.call_soon_threadsafe.call_count == 0 + + def test_livekit_blueprints_are_exposed() -> None: assert teleop_hosted_go2_livekit is not None assert teleop_hosted_xarm6_livekit is not None From de2fcfba9b64557aa06bd880798e5b5588eb68a7 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 7 Aug 2026 20:14:37 -0700 Subject: [PATCH 4/4] fix(teleop): harden LiveKit session lifecycle --- dimos/teleop/hosted/livekit.py | 39 +++++++++++++++++++++++------ dimos/teleop/hosted/test_livekit.py | 32 +++++++++++++++++------ 2 files changed, 56 insertions(+), 15 deletions(-) diff --git a/dimos/teleop/hosted/livekit.py b/dimos/teleop/hosted/livekit.py index 20e4d2ee67..dc732b55c8 100644 --- a/dimos/teleop/hosted/livekit.py +++ b/dimos/teleop/hosted/livekit.py @@ -91,6 +91,7 @@ def __init__(self, **kwargs: Any) -> None: self._video_source: rtc.VideoSource | None = None self._video_track: rtc.LocalVideoTrack | None = None self._video_publish_task: asyncio.Task[None] | None = None + self._room_task: asyncio.Task[None] | None = None self._operator_present = False self._operator_lost = False self._broker = LiveKitBrokerClient(self.config.broker_url, self.config.api_key) @@ -117,7 +118,10 @@ def start(self) -> None: @rpc def stop(self) -> None: self._stop_event.set() - super().stop() + with self._lock: + loop, room_task = self._loop, self._room_task + if loop is not None and room_task is not None and not room_task.done(): + loop.call_soon_threadsafe(room_task.cancel) worker = self._worker if worker is not None: @@ -126,6 +130,7 @@ def stop(self) -> None: logger.error("LiveKit teleop worker did not stop") else: self._worker = None + super().stop() def _register_streams(self) -> None: self.register_disposable(Disposable(self.mux_image.subscribe(self._publish_video))) @@ -154,10 +159,13 @@ def _register_streams(self) -> None: def _run(self) -> None: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) + room_task = loop.create_task(self._run_room()) with self._lock: self._loop = loop + self._room_task = room_task try: - loop.run_until_complete(self._run_room()) + with contextlib.suppress(asyncio.CancelledError): + loop.run_until_complete(room_task) finally: pending = asyncio.all_tasks(loop) for task in pending: @@ -166,21 +174,25 @@ def _run(self) -> None: loop.close() with self._lock: self._loop = None + self._room_task = None async def _run_room(self) -> None: session: LiveKitSession | None = None heartbeat_task: asyncio.Task[None] | None = None + disconnected = asyncio.Event() try: session = await self._broker.create_session( self.config.robot_id, self.config.robot_name, self.config.robot_type ) - await self._connect_room(session) + await self._connect_room(session, disconnected) heartbeat_task = asyncio.create_task( self._broker.heartbeat(session.session_id, self.config.heartbeat_hz) ) self._ready.set() - while not self._stop_event.is_set(): + while not self._stop_event.is_set() and not disconnected.is_set(): await asyncio.sleep(0.1) + if disconnected.is_set() and not self._stop_event.is_set(): + logger.warning("LiveKit room disconnected; ending teleop session") except Exception as error: self._start_error = error self._ready.set() @@ -193,7 +205,7 @@ async def _run_room(self) -> None: if session is not None: await self._broker.close_session(session.session_id) - async def _connect_room(self, session: LiveKitSession) -> None: + async def _connect_room(self, session: LiveKitSession, disconnected: asyncio.Event) -> None: from livekit import rtc room = rtc.Room() @@ -202,7 +214,11 @@ async def _connect_room(self, session: LiveKitSession) -> None: @room.on("data_received") # type: ignore[untyped-decorator] def _data_received(packet: Any) -> None: - self._on_data(getattr(packet, "topic", "") or "", bytes(packet.data)) + self._on_data( + getattr(packet, "topic", "") or "", + bytes(packet.data), + getattr(packet, "participant", None), + ) @room.on("participant_connected") # type: ignore[untyped-decorator] def _participant_connected(participant: Any) -> None: @@ -222,6 +238,7 @@ def _participant_disconnected(participant: Any) -> None: @room.on("disconnected") # type: ignore[untyped-decorator] def _disconnected(_reason: Any) -> None: self._on_operator_lost() + disconnected.set() await room.connect(session.url, session.token) with self._lock: @@ -242,7 +259,11 @@ async def _disconnect_room(self) -> None: def _is_operator(self, participant: Any) -> bool: identity = getattr(participant, "identity", None) - return self.config.operator_identity is None or identity == self.config.operator_identity + if self.config.operator_identity is not None: + return identity == self.config.operator_identity + # The hosted broker mints operator identities as op- and + # viewer identities as viewer--. + return isinstance(identity, str) and identity.startswith("op-") def _has_operator(self, room: rtc.Room, excluded_identity: str | None = None) -> bool: return any( @@ -273,7 +294,9 @@ def _publish_video(self, image: Image) -> None: return loop.call_soon_threadsafe(self._capture_video, width, height, rgba) - def _on_data(self, topic: str, data: bytes) -> None: + def _on_data(self, topic: str, data: bytes, participant: Any = None) -> None: + if not self._is_operator(participant): + return if topic == "state_reliable": self.state_json.publish(data) self.camera_select.publish(data) diff --git a/dimos/teleop/hosted/test_livekit.py b/dimos/teleop/hosted/test_livekit.py index 759efc5902..00923647fa 100644 --- a/dimos/teleop/hosted/test_livekit.py +++ b/dimos/teleop/hosted/test_livekit.py @@ -72,7 +72,9 @@ def module(monkeypatch: pytest.MonkeyPatch) -> LiveKitTestModule: def test_state_packets_fan_out_to_command_and_camera_selection(module: LiveKitTestModule) -> None: - module._on_data("state_reliable", b'{"type":"camera_select"}') + module._on_data( + "state_reliable", b'{"type":"camera_select"}', SimpleNamespace(identity="op-user") + ) module.state_json.publish.assert_called_once_with(b'{"type":"camera_select"}') module.camera_select.publish.assert_called_once_with(b'{"type":"camera_select"}') @@ -84,7 +86,7 @@ def test_command_packets_reach_raw_and_typed_command_ports( command = MagicMock(spec=TwistStamped) mocker.patch.object(TwistStamped, "lcm_decode", return_value=command) - module._on_data("cmd_unreliable", b"command") + module._on_data("cmd_unreliable", b"command", SimpleNamespace(identity="op-user")) module.cmd_raw.publish.assert_called_once_with(b"command") module.cmd_vel_in.publish.assert_called_once_with(command) @@ -95,12 +97,19 @@ def test_non_twist_command_still_reaches_raw_port( ) -> None: mocker.patch.object(TwistStamped, "lcm_decode", side_effect=ValueError("not a twist")) - module._on_data("cmd_unreliable", b"arm-command") + module._on_data("cmd_unreliable", b"arm-command", SimpleNamespace(identity="op-user")) module.cmd_raw.publish.assert_called_once_with(b"arm-command") module.cmd_vel_in.publish.assert_not_called() +def test_viewer_packets_are_ignored(module: LiveKitTestModule) -> None: + module._on_data("cmd_unreliable", b"command", SimpleNamespace(identity="viewer-user-1234")) + + module.cmd_raw.publish.assert_not_called() + module.cmd_vel_in.publish.assert_not_called() + + def test_operator_loss_is_emitted_once(module: LiveKitTestModule) -> None: module._operator_present = True @@ -208,9 +217,12 @@ def register(handler: object) -> object: room.on.side_effect = on mocker.patch("livekit.rtc.Room", return_value=room) + disconnected = asyncio.Event() asyncio.run( - module._connect_room(LiveKitSession("session", "wss://livekit.example", "jwt", "room")) + module._connect_room( + LiveKitSession("session", "wss://livekit.example", "jwt", "room"), disconnected + ) ) assert set(handlers) == { @@ -219,10 +231,16 @@ def register(handler: object) -> object: "participant_disconnected", "disconnected", } - handlers["data_received"](SimpleNamespace(topic="state_reliable", data=b"state")) # type: ignore[operator] - room.remote_participants = {"operator": SimpleNamespace(identity="operator")} + handlers["data_received"]( # type: ignore[operator] + SimpleNamespace( + topic="state_reliable", data=b"state", participant=SimpleNamespace(identity="op-user") + ) + ) + room.remote_participants = {"operator": SimpleNamespace(identity="op-user")} module._operator_present = True - handlers["participant_disconnected"](SimpleNamespace(identity="operator")) # type: ignore[operator] + handlers["participant_disconnected"](SimpleNamespace(identity="op-user")) # type: ignore[operator] + handlers["disconnected"](None) # type: ignore[operator] + assert disconnected.is_set() assert [call.args[0] for call in module.state_json.publish.call_args_list] == [ b"state", b'{"type": "operator_lost"}',