diff --git a/CLAUDE.md b/CLAUDE.md index 42253292..544fd7cb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,6 +37,7 @@ Run the Python tests with `pytest` (config in [pyproject.toml](pyproject.toml); - [record.py](lelab/record.py) — dataset recording (wraps `lerobot.record.record`); patches `lerobot.common.utils.control_utils` keyboard listener so frontend buttons replace arrow-key controls. The worker is a daemon thread, so the server's shutdown hook calls `stop_recording_and_wait()` to let an active session finalize on Ctrl+C. - [teleoperate.py](lelab/teleoperate.py) — leader→follower teleoperation (wraps `lerobot.teleoperate`). +- [replay.py](lelab/replay.py): plays a recorded episode's actions on the follower arm, an in-process take on `lerobot-replay` so the dataset browser can stop it. - [calibrate.py](lelab/calibrate.py) — step-by-step web calibration with a `CalibrationManager` singleton and `_step_complete` threading.Event. - [train.py](lelab/train.py) — wraps the LeRobot training CLI as a subprocess (psutil for lifecycle, queue for log streaming). - [datasets.py](lelab/datasets.py) — dataset listing (local cache + Hub, merged under a `source` field) plus the episode-browsing handlers behind `/dataset-episodes`, `/dataset-episode`, `/dataset-frame`, `/dataset-thumbnails`, `/dataset-motion` and `/dataset-video`. @@ -45,7 +46,7 @@ Run the Python tests with `pytest` (config in [pyproject.toml](pyproject.toml); ### State model -Each feature module owns module-level globals (e.g. `recording_active`, `teleoperation_active`, `current_robot`) protected by threads/locks where needed. There's no shared session object — features are mutually exclusive in practice (you can't teleoperate and record simultaneously) but this is **not** enforced in code. +Each feature module owns module-level globals (e.g. `recording_active`, `teleoperation_active`, `current_robot`) protected by threads/locks where needed. There's no shared session object. Recording, teleoperation, inference and replay all drive the follower, so each start handler checks the others' flags and refuses to run while one is active. Calibration is not part of that check. ### WebSocket broadcast diff --git a/frontend/src/components/dataset/EpisodeViewer.test.tsx b/frontend/src/components/dataset/EpisodeViewer.test.tsx index faff82e4..c8b3f857 100644 --- a/frontend/src/components/dataset/EpisodeViewer.test.tsx +++ b/frontend/src/components/dataset/EpisodeViewer.test.tsx @@ -10,6 +10,7 @@ vi.mock("@/contexts/ApiContext", () => ({ })); vi.mock("./FilmStrip", () => ({ default: () => null })); vi.mock("./MotionTrace", () => ({ default: () => null })); +vi.mock("./ReplayButton", () => ({ default: () => null })); test.each([0, 4])("playback loops at file end with an episode offset of %s seconds", async (offset) => { Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); diff --git a/frontend/src/components/dataset/EpisodeViewer.tsx b/frontend/src/components/dataset/EpisodeViewer.tsx index 83f19b7e..2603ccc9 100644 --- a/frontend/src/components/dataset/EpisodeViewer.tsx +++ b/frontend/src/components/dataset/EpisodeViewer.tsx @@ -15,6 +15,7 @@ import { useApi } from "@/contexts/ApiContext"; import { EpisodeDetail, frameUrl, videoUrl } from "@/lib/datasetApi"; import FilmStrip from "./FilmStrip"; import MotionTrace from "./MotionTrace"; +import ReplayButton from "./ReplayButton"; const RATES = [0.5, 1, 2] as const; type Rate = (typeof RATES)[number]; @@ -160,6 +161,18 @@ const EpisodeViewer: React.FC = ({ repoId, detail, onNavigat } }, []); + // The arm replays at the recording's speed; run the footage alongside it. + const playFromStart = useCallback(() => { + const v = driverRef.current; + if (!v) return; + setRate(1); + seekToFrame(0); + void v.play().then( + () => setPlaying(true), + () => setPlaying(false), + ); + }, [seekToFrame]); + // Keep the element's rate in step with the pills across re-renders and // camera swaps. useEffect(() => { @@ -229,19 +242,26 @@ const EpisodeViewer: React.FC = ({ repoId, detail, onNavigat )} -
- {RATES.map((r) => ( - - ))} +
+ +
+ {RATES.map((r) => ( + + ))} +
diff --git a/frontend/src/components/dataset/ReplayButton.test.tsx b/frontend/src/components/dataset/ReplayButton.test.tsx new file mode 100644 index 00000000..d5d2713e --- /dev/null +++ b/frontend/src/components/dataset/ReplayButton.test.tsx @@ -0,0 +1,134 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot, Root } from "react-dom/client"; +import { afterEach, beforeEach, expect, test, vi } from "vitest"; +import ReplayButton from "./ReplayButton"; +import { RobotRecord } from "@/hooks/useRobots"; + +const backend = vi.hoisted(() => ({ + active: false, + startOk: true, + statusFails: false, + requests: [] as { path: string; body?: unknown }[], +})); +const robot = vi.hoisted(() => ({ current: null as RobotRecord | null })); +const toast = vi.hoisted(() => vi.fn()); +const api = vi.hoisted(() => ({ baseUrl: "http://localhost:8000", fetchWithHeaders: () => {} })); + +vi.mock("@/contexts/ApiContext", () => ({ useApi: () => api })); +vi.mock("@/hooks/useRobots", () => ({ useRobots: () => ({ selectedRecord: robot.current }) })); +vi.mock("@/hooks/use-toast", () => ({ useToast: () => ({ toast }) })); +vi.mock("@/lib/apiClient", () => ({ + apiRequest: async (_baseUrl: string, _fetcher: unknown, path: string, opts?: { body?: unknown }) => { + backend.requests.push({ path, body: opts?.body }); + if (path === "/replay-status") { + if (backend.statusFails) throw new Error("network"); + return { replay_active: backend.active, error: null }; + } + if (path === "/start-replay") { + backend.active = backend.startOk; + return { success: backend.startOk, message: "A replay is already running" }; + } + return { success: true, message: "Replay stopped" }; + }, +})); + +const calibrated: RobotRecord = { + name: "arm", + leader_port: "/dev/leader", + follower_port: "/dev/follower", + leader_config: "leader.json", + follower_config: "follower.json", + cameras: [], + is_clean: true, +}; + +let host: HTMLDivElement; +let root: Root; +const onStart = vi.fn(); + +const button = (label: string) => + [...host.querySelectorAll("button")].find((b) => b.textContent?.includes(label)); + +async function mount() { + await act(async () => { + root.render(); + }); +} + +async function poll() { + await act(async () => { + await vi.advanceTimersByTimeAsync(500); + }); +} + +beforeEach(() => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + vi.useFakeTimers(); + Object.assign(backend, { active: false, startOk: true, statusFails: false, requests: [] }); + robot.current = calibrated; + toast.mockClear(); + onStart.mockClear(); + host = document.createElement("div"); + root = createRoot(host); +}); + +afterEach(() => { + act(() => root.unmount()); + vi.useRealTimers(); +}); + +test("hidden until a calibrated robot is selected", async () => { + robot.current = null; + await mount(); + expect(host.querySelector("button")).toBeNull(); +}); + +test("a replay already running stays stoppable without a selected robot", async () => { + robot.current = null; + backend.active = true; + await mount(); + expect(button("Stop robot")).toBeDefined(); +}); + +test("starting replays this episode and runs the footage alongside", async () => { + await mount(); + await act(async () => button("Play on robot")!.click()); + expect(backend.requests.at(-1)).toMatchObject({ + path: "/start-replay", + body: { dataset_repo_id: "local/ds", episode_index: 2, follower_port: "/dev/follower" }, + }); + expect(onStart).toHaveBeenCalledOnce(); + expect(button("Stop robot")).toBeDefined(); +}); + +test("Stop stays until the status shows the arm released", async () => { + backend.active = true; + await mount(); + await act(async () => button("Stop robot")!.click()); + await poll(); + expect(button("Stop robot")).toBeDefined(); + + backend.active = false; + await poll(); + expect(button("Play on robot")).toBeDefined(); +}); + +test("a failed status poll keeps the Stop button", async () => { + backend.active = true; + await mount(); + backend.statusFails = true; + await poll(); + expect(button("Stop robot")).toBeDefined(); +}); + +test("a refused start reports it and offers Play again", async () => { + backend.startOk = false; + await mount(); + await act(async () => button("Play on robot")!.click()); + expect(onStart).not.toHaveBeenCalled(); + expect(toast).toHaveBeenCalledWith( + expect.objectContaining({ description: "A replay is already running" }), + ); + expect(button("Play on robot")).toBeDefined(); +}); diff --git a/frontend/src/components/dataset/ReplayButton.tsx b/frontend/src/components/dataset/ReplayButton.tsx new file mode 100644 index 00000000..6950b6ff --- /dev/null +++ b/frontend/src/components/dataset/ReplayButton.tsx @@ -0,0 +1,147 @@ +import React, { useEffect, useRef, useState } from "react"; +import { Bot, Loader2, Square } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { useApi } from "@/contexts/ApiContext"; +import { useToast } from "@/hooks/use-toast"; +import { useRobots } from "@/hooks/useRobots"; +import { apiRequest } from "@/lib/apiClient"; + +interface ReplayResponse { + success: boolean; + message: string; +} + +interface ReplayStatus { + replay_active: boolean; + error: string | null; +} + +interface ReplayButtonProps { + repoId: string; + episodeIndex: number; + /** Called once the arm starts moving, so the footage can run alongside it. */ + onStart: () => void; +} + +/** + * Plays the episode's recorded actions on the selected follower arm. Only + * shown when that arm is ready to move. + */ +const ReplayButton: React.FC = ({ repoId, episodeIndex, onStart }) => { + const { baseUrl, fetchWithHeaders } = useApi(); + const { selectedRecord: robot } = useRobots(); + const { toast } = useToast(); + const [phase, setPhase] = useState<"idle" | "starting" | "playing">("idle"); + // Paging to another episode while the arm connects must not drag the + // footage back to the one being replayed. + const episodeRef = useRef(episodeIndex); + episodeRef.current = episodeIndex; + + // A replay keeps running when this button unmounts (another dataset, another + // page), so pick it back up to keep it stoppable. + useEffect(() => { + apiRequest(baseUrl, fetchWithHeaders, "/replay-status") + .then((status) => { + if (status.replay_active) setPhase("playing"); + }) + .catch(() => undefined); + }, [baseUrl, fetchWithHeaders]); + + // Poll only while the arm moves, to notice when the episode ends or fails. + useEffect(() => { + if (phase !== "playing") return; + const timer = setInterval(async () => { + try { + const status = await apiRequest(baseUrl, fetchWithHeaders, "/replay-status"); + if (status.replay_active) return; + setPhase("idle"); + if (status.error) { + toast({ title: "Replay failed", description: status.error, variant: "destructive" }); + } + } catch { + // Keep the Stop button: the next poll retries. + } + }, 500); + return () => clearInterval(timer); + }, [phase, baseUrl, fetchWithHeaders, toast]); + + // Back to Play only once the status poll sees the arm released. + const stop = async () => { + try { + const res = await apiRequest(baseUrl, fetchWithHeaders, "/stop-replay", { + method: "POST", + action: "Stop", + }); + if (!res.success) throw new Error(res.message); + } catch (e) { + toast({ + title: "Could not stop the robot", + description: e instanceof Error ? e.message : String(e), + variant: "destructive", + }); + } + }; + + // A running replay stays stoppable even once the robot is no longer selected. + if (phase === "playing") { + return ( + + ); + } + + if (!robot?.is_clean) return null; + + const start = async () => { + setPhase("starting"); + try { + const res = await apiRequest(baseUrl, fetchWithHeaders, "/start-replay", { + method: "POST", + body: { + follower_port: robot.follower_port, + follower_config: robot.follower_config, + dataset_repo_id: repoId, + episode_index: episodeIndex, + }, + action: "Replay", + }); + if (!res.success) throw new Error(res.message); + setPhase("playing"); + if (episodeRef.current === episodeIndex) onStart(); + } catch (e) { + setPhase("idle"); + toast({ + title: "Could not replay this episode", + description: e instanceof Error ? e.message : String(e), + variant: "destructive", + }); + } + }; + + return ( + + ); +}; + +export default ReplayButton; diff --git a/lelab/record.py b/lelab/record.py index 3d71f161..74044b7b 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -34,7 +34,7 @@ from .camera_config import lerobot_camera_settings, validate_camera_rotations from .episode_media import UnreadableDatasetError, ensure_episode_index from .utils.config import setup_calibration_files, with_lelab_tag -from .utils.devices import safe_disconnect_device +from .utils.devices import follower_lock, safe_disconnect_device logger = logging.getLogger(__name__) @@ -58,9 +58,6 @@ last_recording_info: dict[str, Any] | None = ( None # Snapshot of the most recently completed dataset (for /dataset-info) ) -# Guards the start path so two concurrent POST /start-recording calls cannot -# both pass the active-flag check. -_state_lock = threading.Lock() class RecordingRequest(BaseModel): @@ -251,17 +248,19 @@ def handle_start_recording(request: RecordingRequest) -> dict[str, Any]: phase_start_time, \ last_recording_info - from . import rollout as _rollout, teleoperate as _teleoperate + from . import replay as _replay, rollout as _rollout, teleoperate as _teleoperate # Claim the active flag under the lock so two concurrent starts can't both # pass the precondition check. - with _state_lock: + with follower_lock: if recording_active: return {"success": False, "message": "Recording is already active"} if _teleoperate.teleoperation_active: return {"success": False, "message": "Teleoperation is currently active. Stop it first."} if _rollout.inference_active: return {"success": False, "message": "Inference is currently active. Stop it first."} + if _replay.replay_active: + return {"success": False, "message": "Replay is currently active. Stop it first."} recording_active = True recording_thread = None recording_events = None diff --git a/lelab/replay.py b/lelab/replay.py new file mode 100644 index 00000000..f9d73642 --- /dev/null +++ b/lelab/replay.py @@ -0,0 +1,140 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. + +"""Replay a recorded episode on the follower arm: an in-process take on +lerobot-replay, so the dataset browser can stop it.""" + +import logging +import threading +import time +from typing import Any + +from pydantic import BaseModel + +from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig +from lerobot.utils.robot_utils import precise_sleep + +from . import episode_media +from .utils.config import setup_follower_calibration_file +from .utils.devices import follower_lock, safe_disconnect_device + +logger = logging.getLogger(__name__) + +# Only the worker (or a failed start) clears this, so the follower stays +# claimed until it is actually released. +replay_active = False +replay_thread: threading.Thread | None = None +replay_error: str | None = None +_stop = threading.Event() + + +class ReplayRequest(BaseModel): + follower_port: str + follower_config: str + dataset_repo_id: str + episode_index: int + + +def _episode_actions(repo_id: str, episode_index: int) -> tuple[list[dict[str, float]], int]: + """The episode's recorded actions as follower commands, and the dataset fps.""" + dataset_dir = episode_media.resolve_dataset_dir(repo_id) + info = episode_media.read_info(dataset_dir) + names = info["features"]["action"]["names"] + rows = episode_media.load_episode_actions(dataset_dir, episode_index) + return [dict(zip(names, row, strict=True)) for row in rows], info["fps"] + + +def handle_start_replay(request: ReplayRequest) -> dict[str, Any]: + """Connect the follower synchronously, so a failure reaches the caller, + then send the episode's actions at the recording's rate in the background.""" + global replay_active, replay_thread, replay_error + + from . import record as _record, rollout as _rollout, teleoperate as _teleoperate + + with follower_lock: + if replay_active: + return {"success": False, "message": "A replay is already running"} + if _record.recording_active: + return {"success": False, "message": "Recording is currently active. Stop it first."} + if _teleoperate.teleoperation_active: + return {"success": False, "message": "Teleoperation is currently active. Stop it first."} + if _rollout.inference_active: + return {"success": False, "message": "Inference is currently active. Stop it first."} + replay_active = True + replay_error = None + _stop.clear() + + robot = None + try: + actions, fps = _episode_actions(request.dataset_repo_id, request.episode_index) + + follower_id = setup_follower_calibration_file(request.follower_config) + robot = SO101Follower(SO101FollowerConfig(port=request.follower_port, id=follower_id)) + try: + robot.bus.connect() + except Exception as e: + raise RuntimeError( + f"Could not connect to the follower arm on {request.follower_port}. " + "Make sure it's plugged in and powered on, then try again." + ) from e + robot.bus.write_calibration(robot.calibration) + robot.configure() + except Exception as e: + safe_disconnect_device(robot, logger) + replay_active = False + logger.error(f"Failed to start replay: {e}") + return {"success": False, "message": str(e)} + + def replay_worker(): + global replay_active, replay_error + + try: + for action in actions: + if _stop.is_set(): + break + start = time.perf_counter() + robot.send_action(action) + precise_sleep(1 / fps - (time.perf_counter() - start)) + except Exception as e: + logger.error(f"Replay failed: {e}") + replay_error = str(e) + finally: + safe_disconnect_device(robot, logger) + replay_active = False + logger.info("Replay ended") + + replay_thread = threading.Thread(target=replay_worker, name="replay-worker", daemon=True) + replay_thread.start() + + return {"success": True, "message": "Replay started"} + + +def handle_stop_replay() -> dict[str, Any]: + """Signal the worker and wait for it to release the arm.""" + if not replay_active: + return {"success": False, "message": "No replay is running"} + + _stop.set() + worker = replay_thread + if worker is not None: + worker.join(timeout=5.0) + if worker.is_alive(): + logger.warning("Replay worker did not exit within 5s") + return {"success": False, "message": "The replay has not released the arm yet"} + + return {"success": True, "message": "Replay stopped"} + + +def handle_replay_status() -> dict[str, Any]: + return {"replay_active": replay_active, "error": replay_error} diff --git a/lelab/rollout.py b/lelab/rollout.py index 09c2a8d5..e8038299 100644 --- a/lelab/rollout.py +++ b/lelab/rollout.py @@ -39,6 +39,7 @@ from .camera_config import lerobot_camera_settings, validate_camera_rotations from .utils.config import setup_follower_calibration_file +from .utils.devices import follower_lock logger = logging.getLogger(__name__) @@ -301,10 +302,10 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: global inference_active, _inference_proc, _inference_started_at global _inference_rollout_started_at, _inference_meta - # Mutex with teleop and recording: all three drive the same serial bus. - from . import record as _record, teleoperate as _teleoperate + # Mutex with teleop, recording and replay: they all drive the same serial bus. + from . import record as _record, replay as _replay, teleoperate as _teleoperate - with _state_lock: + with follower_lock, _state_lock: if _teleoperate.teleoperation_active: return { "success": False, @@ -317,6 +318,12 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: "status_code": 409, "message": "Recording is currently active. Stop it first.", } + if _replay.replay_active: + return { + "success": False, + "status_code": 409, + "message": "Replay is currently active. Stop it first.", + } if inference_active: return { "success": False, diff --git a/lelab/server.py b/lelab/server.py index d5fb7372..57943acc 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -62,6 +62,12 @@ handle_stop_recording, handle_upload_dataset, ) +from .replay import ( + ReplayRequest, + handle_replay_status, + handle_start_replay, + handle_stop_replay, +) from .rollout import ( InferenceRequest, handle_inference_status, @@ -354,6 +360,22 @@ def teleoperation_status(): return handle_teleoperation_status() +@app.post("/start-replay") +def start_replay(request: ReplayRequest): + """Replay a recorded episode on the follower arm""" + return handle_start_replay(request) + + +@app.post("/stop-replay") +def stop_replay(): + return handle_stop_replay() + + +@app.get("/replay-status") +def replay_status(): + return handle_replay_status() + + @app.get("/joint-positions") def get_joint_positions(): """Get current robot joint positions""" diff --git a/lelab/teleoperate.py b/lelab/teleoperate.py index c4d17d04..2f998185 100644 --- a/lelab/teleoperate.py +++ b/lelab/teleoperate.py @@ -24,7 +24,7 @@ from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig from .utils.config import setup_calibration_files -from .utils.devices import safe_disconnect_device +from .utils.devices import follower_lock, safe_disconnect_device logger = logging.getLogger(__name__) @@ -33,8 +33,6 @@ teleoperation_thread: threading.Thread | None = None current_robot = None current_teleop = None -# Guards the start path; the worker owns disconnect so stop() does not race. -_state_lock = threading.Lock() class TeleoperateRequest(BaseModel): @@ -118,15 +116,17 @@ def handle_start_teleoperation(request: TeleoperateRequest, websocket_manager=No """ global teleoperation_active, teleoperation_thread, current_robot, current_teleop - from . import record as _record, rollout as _rollout + from . import record as _record, replay as _replay, rollout as _rollout - with _state_lock: + with follower_lock: if teleoperation_active: return {"success": False, "message": "Teleoperation is already active"} if _record.recording_active: return {"success": False, "message": "Recording is currently active. Stop it first."} if _rollout.inference_active: return {"success": False, "message": "Inference is currently active. Stop it first."} + if _replay.replay_active: + return {"success": False, "message": "Replay is currently active. Stop it first."} teleoperation_active = True robot = None diff --git a/lelab/utils/devices.py b/lelab/utils/devices.py index 3b949744..59cda063 100644 --- a/lelab/utils/devices.py +++ b/lelab/utils/devices.py @@ -1,4 +1,4 @@ -"""Device cleanup helpers for LeRobot hardware wrappers. +"""Device ownership and cleanup helpers for LeRobot hardware wrappers. Serial ports are a trust boundary on Windows: if a normal disconnect fails while disabling torque, the COM handle can stay open until the Python process @@ -9,9 +9,14 @@ from __future__ import annotations import logging +import threading from contextlib import suppress from typing import Any +# Recording, teleoperation, inference and replay all drive the follower: each +# start handler checks the others' flags and claims its own under this lock. +follower_lock = threading.Lock() + def safe_disconnect_device(device: Any, logger: logging.Logger, context: str = "cleanup") -> None: """Disconnect a LeRobot device and force-release resources on failure.""" diff --git a/tests/test_replay.py b/tests/test_replay.py new file mode 100644 index 00000000..70f13a78 --- /dev/null +++ b/tests/test_replay.py @@ -0,0 +1,142 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. +"""Tests for lelab.replay: request schema, mutual exclusion with the other +modes, and the start/stop paths that never move the arm.""" + +from __future__ import annotations + +import pytest + + +def _replay_request(): + from lelab.replay import ReplayRequest + + return ReplayRequest( + follower_port="/dev/follower", + follower_config="follower", + dataset_repo_id="local/ds", + episode_index=0, + ) + + +def test_replay_request_requires_an_episode() -> None: + from pydantic import ValidationError + + from lelab.replay import ReplayRequest + + with pytest.raises(ValidationError): + ReplayRequest(follower_port="/dev/follower", follower_config="follower", dataset_repo_id="local/ds") + + +@pytest.mark.parametrize( + ("module", "flag"), + [ + ("record", "recording_active"), + ("teleoperate", "teleoperation_active"), + ("rollout", "inference_active"), + ], +) +def test_replay_refuses_to_start_while_the_arm_is_busy( + module: str, flag: str, monkeypatch: pytest.MonkeyPatch +) -> None: + import importlib + + import lelab.replay as replay + + monkeypatch.setattr(importlib.import_module(f"lelab.{module}"), flag, True) + + result = replay.handle_start_replay(_replay_request()) + assert result["success"] is False + assert "Stop it first" in result["message"] + assert replay.replay_active is False + + +def test_other_modes_refuse_to_start_during_a_replay(monkeypatch: pytest.MonkeyPatch) -> None: + import lelab.replay as replay + from lelab.record import RecordingRequest, handle_start_recording + from lelab.rollout import InferenceRequest, handle_start_inference + from lelab.teleoperate import TeleoperateRequest, handle_start_teleoperation + + monkeypatch.setattr(replay, "replay_active", True) + ports = {"leader_port": "/dev/leader", "follower_port": "/dev/follower"} + configs = {"leader_config": "leader", "follower_config": "follower"} + + results = [ + handle_start_teleoperation(TeleoperateRequest(**ports, **configs)), + handle_start_recording( + RecordingRequest(**ports, **configs, dataset_repo_id="local/ds", single_task="t") + ), + handle_start_inference( + InferenceRequest(follower_port="/dev/follower", follower_config="follower", policy_ref="p") + ), + ] + for result in results: + assert result["success"] is False + assert "Replay" in result["message"] + + +def test_start_replay_reports_connection_failure(monkeypatch: pytest.MonkeyPatch) -> None: + import lelab.replay as replay + + monkeypatch.setattr(replay, "_episode_actions", lambda repo_id, episode: ([{"gripper.pos": 0.0}], 30)) + monkeypatch.setattr(replay, "setup_follower_calibration_file", lambda config: "follower") + + class _Bus: + def connect(self) -> None: + raise RuntimeError("serial port unavailable") + + followers = [] + + class _Follower: + def __init__(self, config) -> None: + self.bus = _Bus() + self.disconnected = False + followers.append(self) + + def disconnect(self) -> None: + self.disconnected = True + + monkeypatch.setattr(replay, "SO101Follower", _Follower) + + result = replay.handle_start_replay(_replay_request()) + + assert result["success"] is False + assert "/dev/follower" in result["message"] + assert followers[0].disconnected is True + assert replay.replay_active is False + + +def test_stop_reports_a_replay_that_still_holds_the_arm(monkeypatch: pytest.MonkeyPatch) -> None: + import threading + + import lelab.replay as replay + + class _StuckWorker: + def join(self, timeout: float) -> None: + pass + + def is_alive(self) -> bool: + return True + + monkeypatch.setattr(replay, "replay_active", True) + monkeypatch.setattr(replay, "replay_thread", _StuckWorker()) + monkeypatch.setattr(replay, "_stop", threading.Event()) + + assert replay.handle_stop_replay()["success"] is False + + +def test_stop_replay_when_idle() -> None: + from lelab.replay import handle_stop_replay + + assert handle_stop_replay()["success"] is False