From 062aa26fd6c19f42d8bb086ffd8130ce92183a00 Mon Sep 17 00:00:00 2001
From: eust-w <longtao.wu@gmail.com>
Date: Wed, 5 Aug 2026 14:29:34 +0800
Subject: [PATCH] Reset recorder state after saving

---
 demonstrations/demo_recorder.py |  3 +-
 tests/test_demo_recorder.py     | 64 +++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_demo_recorder.py

diff --git a/demonstrations/demo_recorder.py b/demonstrations/demo_recorder.py
index 6465e9c..a81a899 100644
--- a/demonstrations/demo_recorder.py
+++ b/demonstrations/demo_recorder.py
@@ -88,7 +88,8 @@ def save_demo(self, filename: Optional[str] = None) -> Optional[Path]:
             self._demo_dir.mkdir(parents=True)
         filename = filename or self._demo.metadata.filename
         filepath = self._demo_dir / filename
-        self._demo.save(filepath)
+        filepath = self._demo.save(filepath)
+        self.stop()
         self._demo = None
         self._env = None
         return filepath
diff --git a/tests/test_demo_recorder.py b/tests/test_demo_recorder.py
new file mode 100644
index 0000000..d0f0f04
--- /dev/null
+++ b/tests/test_demo_recorder.py
@@ -0,0 +1,64 @@
+"""Tests for the DemoRecorder lifecycle."""
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+import demonstrations.demo_recorder as demo_recorder_module
+from demonstrations.demo_recorder import DemoRecorder
+
+
+class _FakeDemo:
+    def __init__(self, *, fail_save=False):
+        self.metadata = SimpleNamespace(filename="demo.safetensors")
+        self.fail_save = fail_save
+
+    def save(self, path):
+        if self.fail_save:
+            raise RuntimeError("simulated save failure")
+        path = Path(path)
+        path.write_bytes(b"demo")
+        return path
+
+
+def test_successful_save_stops_recording_and_allows_next_recording(
+    monkeypatch, tmp_path
+):
+    demos = [_FakeDemo(), _FakeDemo()]
+    monkeypatch.setattr(
+        demo_recorder_module.Demo,
+        "from_env",
+        classmethod(lambda _cls, _env: demos.pop(0)),
+    )
+    recorder = DemoRecorder(tmp_path)
+    environment = object()
+
+    recorder.record(environment)
+    assert recorder.is_recording
+
+    path = recorder.save_demo()
+
+    assert path == tmp_path / "demo.safetensors"
+    assert not recorder.is_recording
+    assert recorder.demo is None
+
+    recorder.record(environment)
+    assert recorder.is_recording
+    assert recorder.demo is not None
+
+
+def test_failed_save_keeps_active_recording_for_retry(monkeypatch, tmp_path):
+    demo = _FakeDemo(fail_save=True)
+    monkeypatch.setattr(
+        demo_recorder_module.Demo,
+        "from_env",
+        classmethod(lambda _cls, _env: demo),
+    )
+    recorder = DemoRecorder(tmp_path)
+    recorder.record(object())
+
+    with pytest.raises(RuntimeError, match="simulated save failure"):
+        recorder.save_demo()
+
+    assert recorder.is_recording
+    assert recorder.demo is demo