File wake_engine.h¶

File List > espos_voice > include > espos_voice > wake_engine.h
/* SPDX-FileCopyrightText: 2026 Dirk Wahrheit */
/* SPDX-License-Identifier: Apache-2.0 */
#pragma once

// WakeEngine — ON-DEVICE wake word via esp-sr AFE + WakeNet.
//
// Runs the wake word detector on the panel itself from the raw mic (no
// network audio), which is far more reliable for far-field / small-MEMS-mic
// audio than streaming to a remote openWakeWord (that path scored the panel's
// mic near zero). esp-sr's AFE gives noise suppression + AGC tuned for exactly
// this embedded case.
//
// The model (e.g. wn9_jarvis_tts) is flashed to the "model" partition and
// loaded at runtime via esp_srmodel_init("model"). Which word is compiled in
// is chosen in sdkconfig (CONFIG_SR_WN_WN9_JARVIS_TTS etc.).
//
// SINGLE MIC CONSUMER: the engine owns the mic while listening. Its feed task
// is the only record_pcm() caller. On a detection it fires a callback; the
// owner (WyomingSatellite) PAUSES the engine (pause()) for the pipeline so
// run_mic() becomes the sole reader, then RESUMES it. pause()/resume() stop
// and restart the feed task cleanly — no concurrent capture.

#include <atomic>
#include <functional>

#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "freertos/task.h"

#include "espos_audio/audio_driver.h"

namespace espos_voice {

class WakeEngine {
 public:
  // on_detect is called (from the fetch task) when the wake word fires.
  // muted_fn (optional) returns true to suppress detection (privacy).
  using DetectFn = std::function<void()>;
  using MutedFn = std::function<bool()>;

  explicit WakeEngine(espos_audio::AudioDriver* audio) : audio_(audio) {}
  ~WakeEngine();

  void set_on_detect(DetectFn fn) { on_detect_ = std::move(fn); }
  void set_muted_fn(MutedFn fn) { muted_fn_ = std::move(fn); }
  // Pre-AFE software gain on the mic feed (1 = off). The onboard MEMS mic is
  // quiet; a small boost can help, but too much just amplifies the noise
  // floor and hurts detection. Set before start().
  void set_input_gain(int g) { input_gain_ = g; }
  // WakeNet detection threshold (0.4-0.9999; 0 = model default). Lower = more
  // sensitive. Set before start().
  void set_threshold(float t) { threshold_ = t; }

  // Initialise AFE + WakeNet from the "model" partition and start the feed +
  // fetch tasks. Returns false if no model is present or AFE init fails
  // (wake is then simply unavailable; the rest of the satellite is fine).
  bool start();
  void stop();

  // Pause / resume mic consumption. pause() blocks until the feed task has
  // released the mic (so the caller can hand it to run_mic()); resume()
  // restarts capture. Idempotent.
  void pause();
  void resume();

  bool running() const { return running_.load(); }
  bool listening() const { return listening_.load(); }  // capturing right now
  uint32_t detections() const { return detections_.load(); }
  const char* word() const { return word_; }

  // Diagnostic: copy the most recent ~2 s of the EXACT PCM fed to WakeNet
  // (post-gain) into `out` (up to max_samples), newest last. Returns the count
  // copied. Lets /mic_probe dump what the detector actually sees on-device.
  size_t pcm_snapshot(int16_t* out, size_t max_samples);
  // Milliseconds since the probe ring was last written, or UINT32_MAX if it
  // never has been. The ring keeps serving its last contents after the feed
  // stops, so an age is the only way to tell live audio from a frozen buffer.
  uint32_t pcm_age_ms() const;
  // Drop any retained probe PCM (privacy: called when the mic is muted). The
  // feed loop stops filling the ring while muted and clears it on this call.
  void clear_probe();

 private:
  static void feed_task_tramp(void* arg);
  static void fetch_task_tramp(void* arg);
  void feed_loop();
  void fetch_loop();
  // Start/stop the mic on the mono or 2ch handle per dual_mic_.
  void capture_start();
  void capture_stop();

  espos_audio::AudioDriver* audio_;
  DetectFn on_detect_;
  MutedFn muted_fn_;

  // esp-sr handles (void* to keep esp-sr headers out of this header).
  void* afe_handle_ = nullptr;  // const esp_afe_sr_iface_t*
  void* afe_data_ = nullptr;    // esp_afe_sr_data_t*
  int feed_chunk_ = 0;          // samples per feed (per channel)
  int feed_channels_ = 1;
  // True when AFE was created with format "MM" and the feed pulls the board's
  // 2-channel [MIC1,MIC2] path (start_capture2/record_pcm2). Decided once at
  // start() from audio_->supports_dual_mic(); the mono path is the fallback.
  bool dual_mic_ = false;
  int input_gain_ = 1;
  float threshold_ = 0.0f;
  char word_[24] = "";

  TaskHandle_t feed_task_ = nullptr;
  TaskHandle_t fetch_task_ = nullptr;
  // Given once by each loop as it returns, so stop() JOINs the tasks before
  // freeing the AFE / releasing the mic (never vTaskDelete mid-inference).
  SemaphoreHandle_t feed_exited_ = nullptr;
  SemaphoreHandle_t fetch_exited_ = nullptr;
  // pause() ⇄ feed_loop() handshake: the feed loop gives this exactly once
  // when it parks *for a pause* (not for mute), and pause() takes it. A
  // generation counter makes each pause wait for a fresh park rather than a
  // stale token.
  SemaphoreHandle_t feed_paused_ = nullptr;

  std::atomic<bool> running_{false};
  std::atomic<bool> paused_{false};
  std::atomic<bool> listening_{false};
  std::atomic<uint32_t> detections_{0};
  // Holds the feed loop out of the AFE while resume() re-arms WakeNet.
  std::atomic<bool> rearming_{false};
  // Published by feed_loop(): true while it is parked (not in record_pcm() or
  // feed()). Lets resume() wait for real quiescence instead of sleeping.
  std::atomic<bool> feed_parked_{false};

  // Diagnostic probe ring: the last ~2 s of post-gain PCM fed to WakeNet.
  static constexpr size_t kProbeSamples = 32000;  // 2 s @ 16 kHz
  int16_t* probe_ = nullptr;
  size_t probe_head_ = 0;
  size_t probe_filled_ = 0;
  // Tick of the last probe write, with a separate written-flag: tick 0 is a
  // legitimate timestamp, so it cannot double as "never written".
  std::atomic<uint32_t> probe_last_write_{0};
  std::atomic<bool> probe_written_{false};
  SemaphoreHandle_t probe_mutex_ = nullptr;
  // Count of probe operations (pcm_snapshot / clear_probe) currently touching
  // probe_mutex_/probe_ from OTHER tasks (the /hello httpd path, the mic-mute
  // widget). stop() waits for this to reach zero before freeing them, so a
  // probe call racing engine teardown can't use freed memory.
  std::atomic<int> probe_busy_{0};
};

}  // namespace espos_voice