feat: post_event/attach_control queue for thread-safe session control

include/trossen_sdk/runtime/session_manager.hpp

@@ -13,6 +13,7 @@
 #include <atomic>
 #include <chrono>
 #include <condition_variable>
+#include <deque>
 #include <filesystem>
 #include <functional>
 #include <memory>
@@ -23,6 +24,7 @@
 #include <vector>
 
 #include "trossen_sdk/hw/producer_base.hpp"
+#include "trossen_sdk/hw/session_control/session_control_capable.hpp"
 #include "trossen_sdk/io/backends/lerobot_v2/lerobot_v2_backend.hpp"
 #include "trossen_sdk/io/backends/trossen_mcap/trossen_mcap_backend.hpp"
 #include "trossen_sdk/io/sink.hpp"
@@ -197,6 +199,33 @@ class SessionManager {
    */
   void request_rerecord();
 
+  /**
+   * @brief Post a session-control event from any thread.
+   *
+   * Thread-safe. The event is queued and applied on the orchestrating thread
+   * by monitor_episode() / wait_for_reset(), so a hardware callback never
+   * touches the single-threaded session state directly.
+   */
+  void post_event(hw::session_control::SessionControlEvent event);
+
+  /**
+   * @brief Attach a push-based session-control source (VR buttons, keyboard, …).
+   *
+   * Installs callbacks that forward intents through post_event(), starts the
+   * source, and remembers it so shutdown() stops it (joining its thread)
+   * before the session is torn down. Every source drives the session through
+   * the same thread-safe queue, so keyboard and VR are interchangeable.
+   *
+   * @note The manager holds a shared_ptr to `source`, so the source is
+   * guaranteed to outlive the SessionManager's use of it (until shutdown() /
+   * destruction releases it). This removes the dangling-pointer risk a
+   * non-owning pointer would carry if the caller dropped the source early.
+   * @note Events are consumed by the synchronous monitor_episode() /
+   * wait_for_reset() loop, so drive the session with those rather than
+   * start_async_monitoring().
+   */
+  void attach_control(std::shared_ptr<hw::session_control::SessionControlCapable> source);
+
   /**
    * @brief Check if an episode is currently running
    *
@@ -517,6 +546,34 @@ class SessionManager {
   /// @brief Flag indicating that re-record was requested
   std::atomic<bool> rerecord_requested_{false};
 
+  /// @brief Phase in which queued control events are interpreted.
+  enum class ControlPhase { kRecording, kReset };
+
+  /// @brief Apply queued session-control events on the orchestrating thread,
+  /// translating each into the loop's control signals for the given phase.
+  void drain_control_events(ControlPhase phase);
+
+  // Grants the control-integration unit tests access to drain_control_events()
+  // and the resulting control flags.
+  friend class SessionManagerControlTest;
+
+  /// @brief Thread-safe queue of session-control events posted by sources.
+  std::deque<hw::session_control::SessionControlEvent> control_events_;
+
+  /// @brief Guards control_events_.
+  std::mutex control_events_mutex_;
+
+  /// @brief Attached control sources; stopped in shutdown() before teardown.
+  /// Held by shared_ptr so a source cannot be destroyed while the manager may
+  /// still stop it.
+  std::vector<std::shared_ptr<hw::session_control::SessionControlCapable>> control_sources_;
+
+  /// @brief Set by drain when a control event asks the current episode to stop.
+  std::atomic<bool> episode_stop_requested_{false};
+
+  /// @brief Set by drain when kStart during recording asks to skip the reset.
+  std::atomic<bool> skip_reset_{false};
+
   /// @brief Flag indicating that episode final statistics were emitted
   mutable std::atomic<bool> stats_emitted_{false};
 

src/runtime/session_manager.cpp

@@ -93,6 +93,13 @@ SessionManager::~SessionManager() {
 }
 
 void SessionManager::shutdown() {
+  // Stop attached control sources first (joins their reader threads) so no
+  // queued event can drive the session while it is being torn down.
+  for (auto& source : control_sources_) {
+    if (source) source->stop();
+  }
+  control_sources_.clear();
+
   // stop_episode() runs first so episode-ended callbacks fire against a still-non-
   // terminal SessionManager; shutdown_called_ latches afterwards and gates start_episode().
   stop_episode();
@@ -891,19 +898,35 @@ UserAction SessionManager::wait_for_reset() {
   reset_signaled_.store(false);
   rerecord_requested_.store(false);
 
+  // A kStart received while recording asked to skip the reset wait entirely.
+  if (skip_reset_.exchange(false)) {
+    return trossen::utils::g_stop_requested ? UserAction::kStop
+                                            : UserAction::kContinue;
+  }
+
   // Enable raw terminal input to detect keypresses without Enter
   trossen::utils::RawModeGuard raw_mode;
 
   // Helper: wait up to 100ms, waking early on signal_reset_complete()
   auto wait_or_signal = [this]() {
     std::unique_lock<std::mutex> lk(reset_mutex_);
     reset_cv_.wait_for(lk, std::chrono::milliseconds(100), [this]() {
-      return reset_signaled_.load() || rerecord_requested_.load() ||
-             trossen::utils::g_stop_requested;
+      if (reset_signaled_.load() || rerecord_requested_.load() ||
+          trossen::utils::g_stop_requested) {
+        return true;
+      }
+      // Also wake when a control-source event is queued so poll_keys() drains
+      // it promptly instead of waiting out the interval.
+      std::lock_guard<std::mutex> qlk(control_events_mutex_);
+      return !control_events_.empty();
     });
   };
 
   auto poll_keys = [this]() -> UserAction {
+    // Apply any queued control-source events (VR buttons, etc.) on this thread.
+    drain_control_events(ControlPhase::kReset);
+    if (rerecord_requested_.load()) return UserAction::kReRecord;
+
     auto key = trossen::utils::poll_keypress();
     if (key == trossen::utils::KeyPress::kRightArrow) {
       reset_signaled_.store(true);
@@ -968,24 +991,108 @@ void SessionManager::signal_reset_complete() {
   reset_cv_.notify_all();
 }
 
+void SessionManager::post_event(hw::session_control::SessionControlEvent event) {
+  using Event = hw::session_control::SessionControlEvent;
+  if (event == Event::kNone) return;
+  {
+    std::lock_guard<std::mutex> lock(control_events_mutex_);
+    control_events_.push_back(event);
+  }
+  // Wake wait_for_reset() if it is blocked so the queued event is drained on
+  // the next poll pass instead of waiting out the ~100 ms interval.
+  reset_cv_.notify_all();
+}
+
+void SessionManager::attach_control(
+  std::shared_ptr<hw::session_control::SessionControlCapable> source)
+{
+  if (!source) return;
+  using Event = hw::session_control::SessionControlEvent;
+  source->set_callbacks(
+    [this](Event e) { post_event(e); },
+    [this] {
+      std::cerr << "\n[session] control source disconnected; ending session."
+                << std::endl;
+      post_event(Event::kStopSession);
+    });
+  source->start();
+  control_sources_.push_back(std::move(source));
+}
+
+void SessionManager::drain_control_events(ControlPhase phase) {
+  using Event = hw::session_control::SessionControlEvent;
+
+  std::deque<Event> events;
+  {
+    std::lock_guard<std::mutex> lock(control_events_mutex_);
+    events.swap(control_events_);
+  }
+
+  for (const Event e : events) {
+    switch (e) {
+      case Event::kStart:
+        // "Advance": during recording, end the episode and skip the reset
+        // wait; during reset, resume immediately.
+        if (phase == ControlPhase::kRecording) {
+          skip_reset_.store(true);
+          episode_stop_requested_.store(true);
+        } else {
+          reset_signaled_.store(true);
+        }
+        break;
+      case Event::kStopEarly:
+        // Stop the current episode (a normal reset follows). Only meaningful
+        // while recording.
+        if (phase == ControlPhase::kRecording) {
+          episode_stop_requested_.store(true);
+        }
+        break;
+      case Event::kRerecord:
+        rerecord_requested_.store(true);
+        break;
+      case Event::kStopSession:
+        trossen::utils::g_stop_requested = true;
+        break;
+      case Event::kNone:
+        break;
+    }
+  }
+}
+
 UserAction SessionManager::monitor_episode(
   std::chrono::duration<double> update_interval,
   std::chrono::duration<double> sleep_interval,
   bool print_stats)
 {
   rerecord_requested_.store(false);
+  episode_stop_requested_.store(false);
 
   // Enable raw terminal input to detect keypresses during recording
   trossen::utils::RawModeGuard raw_mode;
 
   auto last_update = std::chrono::steady_clock::now();
 
   while (!are_final_stats_emitted()) {
+    // Apply any queued control-source events (VR buttons, etc.) on this thread.
+    drain_control_events(ControlPhase::kRecording);
+
+    // End-session takes priority over any episode-level action queued in the
+    // same drain.
+    if (trossen::utils::g_stop_requested) {
+      return UserAction::kStop;
+    }
+
     // Check for re-record request
     if (rerecord_requested_.load()) {
       return UserAction::kReRecord;
     }
 
+    // A control source asked to end the episode (kStart advances, kStopEarly
+    // just stops); return so the main loop stops it on this thread.
+    if (episode_stop_requested_.load()) {
+      return UserAction::kContinue;
+    }
+
     // Poll for arrow keys: left = re-record, right = early exit to next episode
     auto key = trossen::utils::poll_keypress();
     if (key == trossen::utils::KeyPress::kLeftArrow) {

tests/CMakeLists.txt

@@ -65,6 +65,9 @@ target_link_libraries(test_session_manager PRIVATE trossen_sdk GTest::gtest_main
 add_executable(test_session_manager_discard test_session_manager_discard.cpp)
 target_link_libraries(test_session_manager_discard PRIVATE trossen_sdk GTest::gtest_main)
 
+add_executable(test_session_manager_control test_session_manager_control.cpp)
+target_link_libraries(test_session_manager_control PRIVATE trossen_sdk GTest::gtest_main)
+
 add_executable(test_keyboard_input_utils test_keyboard_input_utils.cpp)
 target_link_libraries(test_keyboard_input_utils PRIVATE trossen_sdk GTest::gtest_main)
 
@@ -119,6 +122,7 @@ gtest_discover_tests(test_sink)
 gtest_discover_tests(test_scheduler)
 gtest_discover_tests(test_session_manager)
 gtest_discover_tests(test_session_manager_discard)
+gtest_discover_tests(test_session_manager_control)
 gtest_discover_tests(test_keyboard_input_utils)
 gtest_discover_tests(test_announce)
 gtest_discover_tests(test_teleop_space_helpers)