feat: Wave 3 — plugin system + custom wake-word trainer

#29 Plugin system:
- New jarvis-core::plugins module: discovers user packs in
  %APPDATA%\com.priler.jarvis\plugins\<name>\command.toml so authors can
  ship voice commands without rebuilding. Sandbox capped at "standard" —
  plugins cannot escalate to "full" (which would expose `os`). Disabled
  via a `disabled` flag file. Malformed packs warn and skip; never poison
  the rest of the list. 8 unit tests.
- commands::parse_commands() merges plugins into the loaded list.
- New /plugins GUI page with enable/disable switches, error reporting,
  "Open folder" button. Tauri commands plugins_list / plugins_set_enabled
  / plugins_open_folder. Header gets a "Плагины" button. i18n keys added
  for ru/en/ua.

#8 Custom wake-word trainer wizard:
- New jarvis-core::wake_trainer module: opens its own pv_recorder
  instance, records N short PCM clips, WAV-encodes them in memory, then
  calls rustpotter's WakewordRef::new_from_sample_buffers to train and
  persist a .rpw model under %APPDATA%/com.priler.jarvis/wake_words/.
  Keepsake WAVs are also dumped for retraining later. Sanitises names to
  block path traversal. 5 unit tests.
- Settings gain `custom_wake_word: String`; listener/rustpotter.rs now
  loads the user's selected model first and always falls back to the
  bundled default so the assistant keeps working even if the custom
  file is missing.
- New /wake-trainer GUI page: stepper UI for record-sample → train,
  shows recorded count, threshold slider, refuses to start while
  jarvis-app is running (mic exclusivity). Lists existing trained
  models with size + delete button.
- 8 Tauri commands wired through (status/defaults/start/record_sample/
  finish/cancel/list_models/delete_model).

Tests: 115 → 128 (+8 plugins +5 wake_trainer). Release builds green for
all three binaries (jarvis-app / jarvis-cli / jarvis-gui).
This commit is contained in:
Bossiara13 2026-05-16 13:26:42 +03:00
parent c4b22618f8
commit 4e21024509
17 changed files with 1625 additions and 24 deletions

View file

@ -9,7 +9,7 @@ use seqdiff::ratio;
mod structs;
pub use structs::*;
use crate::{config, i18n, APP_DIR};
use crate::{config, i18n, plugins, APP_DIR};
#[cfg(feature = "lua")]
use crate::lua::{self, SandboxLevel, CommandContext};
@ -128,10 +128,22 @@ pub fn parse_commands() -> Result<Vec<JCommandsList>, String> {
});
}
// Merge plugin packs from <APP_CONFIG_DIR>/plugins/. Plugin loading is
// best-effort: malformed or missing plugin dirs are logged and skipped, never
// blocking startup. See `plugins.rs` for the layout & sandbox policy.
let plugin_packs = plugins::discover();
let plugin_count = plugin_packs.len();
commands.extend(plugin_packs);
if commands.is_empty() {
Err("No commands found".into())
} else {
info!("Loaded {} command pack(s)", commands.len());
info!(
"Loaded {} command pack(s) total ({} built-in + {} plugin)",
commands.len(),
commands.len() - plugin_count,
plugin_count
);
Ok(commands)
}
}

View file

@ -46,6 +46,12 @@ pub struct Settings {
/// auto-detect.
#[serde(default)]
pub tts_backend: String,
/// Name of a custom Rustpotter wake-word model (`.rpw` stem) trained by the
/// user via the wake-word wizard. Empty string means "use the bundled
/// default model". Looked up under `<APP_CONFIG_DIR>/wake_words/<name>.rpw`.
#[serde(default)]
pub custom_wake_word: String,
}
fn default_intent_backend() -> String { config::DEFAULT_INTENT_BACKEND.to_string() }
@ -75,6 +81,7 @@ impl Settings {
"api_key__openai" => Some(self.api_keys.openai.clone()),
"llm_backend" => Some(self.llm_backend.clone()),
"tts_backend" => Some(self.tts_backend.clone()),
"custom_wake_word" => Some(self.custom_wake_word.clone()),
_ => None,
}
}
@ -151,6 +158,13 @@ impl Settings {
other => return Err(format!("unknown tts_backend: '{}'", other)),
}
}
"custom_wake_word" => {
// empty = bundled default; otherwise treated as a file-stem
// under <APP_CONFIG_DIR>/wake_words/<name>.rpw. We accept it
// verbatim — the listener will fall back to the default if the
// file isn't there.
self.custom_wake_word = val.trim().to_string();
}
_ => return Err(format!("unknown setting: '{}'", key)),
}
Ok(())
@ -175,6 +189,7 @@ impl Settings {
"api_key__openai",
"llm_backend",
"tts_backend",
"custom_wake_word",
]
}
}
@ -209,6 +224,7 @@ impl Default for Settings {
llm_backend: String::new(),
tts_backend: String::new(),
custom_wake_word: String::new(),
}
}
}

View file

@ -162,4 +162,5 @@ settings-profile = Profile
# Header buttons
header-macros = Macros
header-scheduler = Schedule
header-memory = Memory
header-memory = Memory
header-plugins = Plugins

View file

@ -162,4 +162,5 @@ settings-profile = Профиль
# Header buttons
header-macros = Макросы
header-scheduler = Расписание
header-memory = Память
header-memory = Память
header-plugins = Плагины

View file

@ -162,4 +162,5 @@ settings-profile = Профіль
# Header buttons
header-macros = Макроси
header-scheduler = Розклад
header-memory = Пам'ять
header-memory = Пам'ять
header-plugins = Плагіни

View file

@ -62,6 +62,10 @@ pub mod scheduler;
pub mod macros;
pub mod plugins;
pub mod wake_trainer;
#[cfg(feature = "llm")]
pub mod llm;

View file

@ -3,7 +3,7 @@ use std::sync::Mutex;
use once_cell::sync::OnceCell;
use rustpotter::Rustpotter;
use crate::config;
use crate::{config, APP_CONFIG_DIR, DB};
// store rustpotter instance
static RUSTPOTTER: OnceCell<Mutex<Rustpotter>> = OnceCell::new();
@ -14,26 +14,56 @@ pub fn init() -> Result<(), ()> {
// create rustpotter instance
match Rustpotter::new(&rustpotter_config) {
Ok(mut rinstance) => {
// success
// wake word files list
// @TODO. Make it configurable via GUI for custom user voice.
let rustpotter_wake_word_files: [&str; 1] = [
"resources/rustpotter/jarvis-default.rpw",
// "rustpotter/jarvis-community-1.rpw",
// "rustpotter/jarvis-community-2.rpw",
// "rustpotter/jarvis-community-3.rpw",
// "rustpotter/jarvis-community-4.rpw",
// "rustpotter/jarvis-community-5.rpw",
];
// Resolve wake-word file list. Priority:
// 1. user-trained custom model (settings.custom_wake_word)
// 2. bundled default
// We always try the bundled default last so the assistant keeps
// working even if the custom model is missing on disk.
let mut loaded_any = false;
// load wake word files
for rpw in rustpotter_wake_word_files {
// @TODO: Change wakeword key to something else?
if let Err(e) = rinstance.add_wakeword_from_file(rpw, rpw) {
error!("Failed to load wakeword file '{}': {}", rpw, e);
if let Some(db) = DB.get() {
let stem = db.read().custom_wake_word.clone();
if !stem.is_empty() {
if let Some(cfg_dir) = APP_CONFIG_DIR.get() {
let custom = cfg_dir
.join("wake_words")
.join(format!("{}.rpw", stem));
if custom.is_file() {
let path_str = custom.to_string_lossy().to_string();
match rinstance.add_wakeword_from_file(&path_str, &path_str) {
Ok(_) => {
info!("Loaded custom wake-word: {}", path_str);
loaded_any = true;
}
Err(e) => warn!(
"Failed to load custom wakeword '{}': {}",
path_str, e
),
}
} else {
warn!(
"Custom wake-word '{}' selected but file missing: {}",
stem,
custom.display()
);
}
}
}
}
const DEFAULT: &str = "resources/rustpotter/jarvis-default.rpw";
if let Err(e) = rinstance.add_wakeword_from_file(DEFAULT, DEFAULT) {
if !loaded_any {
error!("Failed to load default wakeword '{}': {}", DEFAULT, e);
}
} else {
loaded_any = true;
}
if !loaded_any {
error!("No wake-word models loaded; Rustpotter will not detect anything.");
}
// store
let _ = RUSTPOTTER.set(Mutex::new(rinstance));
}

View file

@ -0,0 +1,379 @@
//! Plugin system — user-installed command packs that live outside the read-only
//! `resources/commands/` tree.
//!
//! Why: shipping new voice packs as built-ins requires a rebuild + release. End
//! users (and external authors) want to drop a folder into a writable location
//! and have Jarvis pick it up. That's exactly what this module enables.
//!
//! Layout (per pack):
//!
//! ```text
//! <APP_CONFIG_DIR>/plugins/
//! <pack_name>/
//! command.toml (required — same schema as built-in packs)
//! <script>.lua (zero or more scripts referenced from toml)
//! disabled (optional empty file — if present, pack is skipped)
//! ```
//!
//! On Windows this resolves to `%APPDATA%\com.priler.jarvis\plugins\`.
//!
//! Security: plugin Lua scripts run under the same sandbox the `command.toml`
//! declares, but we cap the maximum sandbox level at `Standard`. Plugin authors
//! cannot reach `Full` (which gives `os` access). This is the only place where
//! a plugin pack differs from a built-in pack.
//!
//! Discovery is read-only and tolerant: a malformed `command.toml` logs a
//! warning and is skipped; it never poisons the whole list.
//!
//! Reload: `commands::reload_global()` rescans both built-in and plugin packs
//! and atomically swaps `COMMANDS_LIST`. The jarvis-app FS watcher fires this
//! on any change under `plugins/`.
use std::fs;
use std::path::{Path, PathBuf};
use crate::commands::JCommandsList;
use crate::APP_CONFIG_DIR;
const PLUGINS_DIR_NAME: &str = "plugins";
const DISABLED_FLAG: &str = "disabled";
const MAX_SANDBOX_FOR_PLUGINS: &str = "standard";
/// Metadata about a single plugin pack on disk.
#[derive(Debug, Clone, serde::Serialize)]
pub struct PluginInfo {
pub name: String,
pub path: PathBuf,
pub enabled: bool,
pub command_count: usize,
pub error: Option<String>,
}
/// Returns the on-disk plugins directory (`<APP_CONFIG_DIR>/plugins`).
/// Creates it on first call so the user can find and populate it.
pub fn plugins_dir() -> Result<PathBuf, String> {
let cfg = APP_CONFIG_DIR
.get()
.ok_or_else(|| "APP_CONFIG_DIR not initialised".to_string())?;
let dir = cfg.join(PLUGINS_DIR_NAME);
if !dir.exists() {
fs::create_dir_all(&dir).map_err(|e| format!("create plugins dir: {}", e))?;
}
Ok(dir)
}
/// Walk the plugins directory and parse every `command.toml`.
/// Returns parsed packs; malformed packs are logged and skipped.
pub fn discover() -> Vec<JCommandsList> {
let dir = match plugins_dir() {
Ok(d) => d,
Err(e) => {
warn!("Plugins discovery skipped: {}", e);
return Vec::new();
}
};
discover_in(&dir)
}
/// Same as `discover()` but takes an explicit directory — used by tests.
pub fn discover_in(dir: &Path) -> Vec<JCommandsList> {
let mut out = Vec::new();
let read = match fs::read_dir(dir) {
Ok(r) => r,
Err(e) => {
// not a hard error — directory may not exist yet
debug!("Plugins dir unreadable {}: {}", dir.display(), e);
return out;
}
};
for entry in read.flatten() {
let pack_path = entry.path();
if !pack_path.is_dir() {
continue;
}
if pack_path.join(DISABLED_FLAG).is_file() {
debug!("Plugin '{}' is disabled, skipping", pack_path.display());
continue;
}
let toml_file = pack_path.join("command.toml");
if !toml_file.is_file() {
continue;
}
match parse_one(&pack_path, &toml_file) {
Ok(list) => out.push(list),
Err(e) => warn!("Plugin '{}' load failed: {}", pack_path.display(), e),
}
}
if !out.is_empty() {
info!(
"Loaded {} plugin pack(s) from {}",
out.len(),
dir.display()
);
}
out
}
fn parse_one(pack_path: &Path, toml_file: &Path) -> Result<JCommandsList, String> {
let body = fs::read_to_string(toml_file).map_err(|e| format!("read: {}", e))?;
let mut pack: JCommandsList =
toml::from_str(&body).map_err(|e| format!("parse: {}", e))?;
// enforce sandbox cap and validate script references
for cmd in &mut pack.commands {
if cmd.cmd_type == "lua" {
let requested = if cmd.sandbox.is_empty() {
MAX_SANDBOX_FOR_PLUGINS
} else {
cmd.sandbox.as_str()
};
if requested == "full" {
warn!(
"Plugin '{}': command '{}' requested 'full' sandbox; \
downgrading to 'standard' (plugin policy)",
pack_path.display(),
cmd.id
);
cmd.sandbox = MAX_SANDBOX_FOR_PLUGINS.to_string();
} else if cmd.sandbox.is_empty() {
cmd.sandbox = MAX_SANDBOX_FOR_PLUGINS.to_string();
}
if !cmd.script.is_empty() {
let s = pack_path.join(&cmd.script);
if !s.is_file() {
return Err(format!(
"command '{}' references missing script: {}",
cmd.id,
s.display()
));
}
}
}
}
pack.path = pack_path.to_path_buf();
Ok(pack)
}
/// Lightweight listing for the GUI. Reports each pack regardless of enable state.
pub fn list() -> Vec<PluginInfo> {
let dir = match plugins_dir() {
Ok(d) => d,
Err(_) => return Vec::new(),
};
list_in(&dir)
}
pub fn list_in(dir: &Path) -> Vec<PluginInfo> {
let mut out = Vec::new();
let read = match fs::read_dir(dir) {
Ok(r) => r,
Err(_) => return out,
};
for entry in read.flatten() {
let pack_path = entry.path();
if !pack_path.is_dir() {
continue;
}
let toml_file = pack_path.join("command.toml");
if !toml_file.is_file() {
continue;
}
let name = pack_path
.file_name()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_default();
let enabled = !pack_path.join(DISABLED_FLAG).is_file();
let (count, error) = match parse_one(&pack_path, &toml_file) {
Ok(p) => (p.commands.len(), None),
Err(e) => (0, Some(e)),
};
out.push(PluginInfo {
name,
path: pack_path,
enabled,
command_count: count,
error,
});
}
out.sort_by(|a, b| a.name.cmp(&b.name));
out
}
/// Touch `disabled` flag inside the pack folder so the next discovery skips it.
pub fn set_enabled(name: &str, enabled: bool) -> Result<(), String> {
let dir = plugins_dir()?;
let pack = dir.join(name);
if !pack.is_dir() {
return Err(format!("plugin '{}' not found", name));
}
let flag = pack.join(DISABLED_FLAG);
if enabled {
if flag.is_file() {
fs::remove_file(&flag).map_err(|e| format!("remove flag: {}", e))?;
}
} else if !flag.is_file() {
fs::write(&flag, b"").map_err(|e| format!("write flag: {}", e))?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn write(path: &Path, content: &str) {
if let Some(p) = path.parent() {
std::fs::create_dir_all(p).unwrap();
}
std::fs::write(path, content).unwrap();
}
fn good_toml() -> &'static str {
r#"
[[commands]]
id = "my_plugin.hello"
type = "lua"
script = "hello.lua"
sandbox = "standard"
timeout = 3000
[commands.phrases]
ru = ["привет плагин"]
"#
}
#[test]
fn discovers_packs_with_command_toml() {
let dir = TempDir::new().unwrap();
let pack = dir.path().join("greeter");
write(&pack.join("command.toml"), good_toml());
write(&pack.join("hello.lua"), "return jarvis.cmd.ok('hi')");
let packs = discover_in(dir.path());
assert_eq!(packs.len(), 1, "expected exactly one pack");
assert_eq!(packs[0].commands.len(), 1);
assert_eq!(packs[0].commands[0].id, "my_plugin.hello");
assert_eq!(packs[0].path, pack);
}
#[test]
fn disabled_flag_skips_pack() {
let dir = TempDir::new().unwrap();
let pack = dir.path().join("greeter");
write(&pack.join("command.toml"), good_toml());
write(&pack.join("hello.lua"), "return jarvis.cmd.ok('hi')");
write(&pack.join("disabled"), "");
let packs = discover_in(dir.path());
assert!(packs.is_empty(), "disabled pack must be skipped");
}
#[test]
fn missing_script_is_an_error() {
let dir = TempDir::new().unwrap();
let pack = dir.path().join("broken");
write(&pack.join("command.toml"), good_toml());
// intentionally no hello.lua
let packs = discover_in(dir.path());
assert!(packs.is_empty(), "pack with missing script must not load");
}
#[test]
fn malformed_toml_is_skipped_not_fatal() {
let dir = TempDir::new().unwrap();
let pack_bad = dir.path().join("bad");
write(&pack_bad.join("command.toml"), "this is not valid toml ::::");
let pack_good = dir.path().join("good");
write(&pack_good.join("command.toml"), good_toml());
write(&pack_good.join("hello.lua"), "return jarvis.cmd.ok('hi')");
let packs = discover_in(dir.path());
assert_eq!(packs.len(), 1, "good pack should still load despite bad sibling");
}
#[test]
fn full_sandbox_is_downgraded_to_standard() {
let dir = TempDir::new().unwrap();
let pack = dir.path().join("evil");
let toml = r#"
[[commands]]
id = "evil.run"
type = "lua"
script = "evil.lua"
sandbox = "full"
timeout = 1000
[commands.phrases]
ru = ["взломай систему"]
"#;
write(&pack.join("command.toml"), toml);
write(&pack.join("evil.lua"), "return jarvis.cmd.ok('nope')");
let packs = discover_in(dir.path());
assert_eq!(packs.len(), 1);
assert_eq!(
packs[0].commands[0].sandbox, "standard",
"plugin must not be allowed to escalate to full sandbox"
);
}
#[test]
fn empty_sandbox_defaults_to_standard() {
let dir = TempDir::new().unwrap();
let pack = dir.path().join("plain");
let toml = r#"
[[commands]]
id = "plain.go"
type = "lua"
script = "go.lua"
timeout = 1000
[commands.phrases]
ru = ["иди"]
"#;
write(&pack.join("command.toml"), toml);
write(&pack.join("go.lua"), "return jarvis.cmd.ok('ok')");
let packs = discover_in(dir.path());
assert_eq!(packs.len(), 1);
assert_eq!(packs[0].commands[0].sandbox, "standard");
}
#[test]
fn list_reports_enabled_and_disabled() {
let dir = TempDir::new().unwrap();
let a = dir.path().join("alpha");
write(&a.join("command.toml"), good_toml());
write(&a.join("hello.lua"), "return jarvis.cmd.ok('hi')");
let b = dir.path().join("bravo");
write(&b.join("command.toml"), good_toml());
write(&b.join("hello.lua"), "return jarvis.cmd.ok('hi')");
write(&b.join("disabled"), "");
let listing = list_in(dir.path());
assert_eq!(listing.len(), 2);
let alpha = listing.iter().find(|p| p.name == "alpha").unwrap();
let bravo = listing.iter().find(|p| p.name == "bravo").unwrap();
assert!(alpha.enabled);
assert!(!bravo.enabled);
assert_eq!(alpha.command_count, 1);
}
#[test]
fn non_directories_are_ignored() {
let dir = TempDir::new().unwrap();
write(&dir.path().join("README.md"), "# hi");
write(&dir.path().join("strayfile.toml"), good_toml());
let packs = discover_in(dir.path());
assert!(packs.is_empty());
}
}

View file

@ -0,0 +1,439 @@
//! Custom wake-word trainer — guides the user through recording a handful of
//! samples of their chosen wake phrase, then trains a Rustpotter `.rpw` model
//! that can replace the bundled `jarvis-default.rpw`.
//!
//! Flow (driven by `jarvis-gui` over Tauri commands):
//!
//! ```text
//! start(name, target=10)
//! -> session entered; mic opened (exclusive); state = Recording { ... }
//! record_sample(seconds=1.5) ×N
//! -> blocks for `seconds`, captures i16 samples, appends to buffer
//! -> WAV-encodes each clip in memory so rustpotter can consume it
//! finish(threshold=0.5)
//! -> calls rustpotter's WakewordRef::new_from_sample_buffers
//! -> serialises .rpw into <APP_CONFIG_DIR>/wake_words/<name>.rpw
//! -> mic released; state = Idle
//! cancel()
//! -> mic released; state = Idle, samples discarded
//! ```
//!
//! Concurrency: the trainer holds its own `pv_recorder` instance so it doesn't
//! collide with the `recorder` module that jarvis-app uses. pv_recorder DOES
//! claim the device exclusively, so the trainer refuses to start if jarvis-app
//! is recording.
//!
//! Output layout:
//! - `<APP_CONFIG_DIR>/wake_words/<name>.rpw` — the trained model
//! - `<APP_CONFIG_DIR>/wake_words/<name>/sample_NN.wav` — keepsake samples,
//! useful if the user wants to retrain with different parameters later
//!
//! Selecting which `.rpw` is active is handled by the settings page (radio
//! between bundled default and any custom models discovered here).
use once_cell::sync::OnceCell;
use parking_lot::Mutex;
use rustpotter::{WakewordRef, WakewordRefBuildFromBuffers, WakewordSave};
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::time::{Duration, Instant};
use crate::APP_CONFIG_DIR;
const WAKE_WORDS_DIR_NAME: &str = "wake_words";
const SAMPLE_RATE: u32 = 16_000;
const FRAME_LEN: i32 = 512;
const MFCC_SIZE: u16 = 16;
pub const DEFAULT_SAMPLE_SECONDS: f32 = 1.5;
pub const DEFAULT_TARGET_SAMPLES: u8 = 10;
pub const MIN_TARGET_SAMPLES: u8 = 5;
pub const MAX_TARGET_SAMPLES: u8 = 30;
#[derive(Debug, Clone, serde::Serialize)]
pub struct TrainerStatus {
pub recording: bool,
pub session_name: Option<String>,
pub target_samples: u8,
pub collected: usize,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct TrainedModelInfo {
pub name: String,
pub path: String,
pub size_bytes: u64,
pub modified: i64,
}
struct Session {
name: String,
target_samples: u8,
// each sample is a WAV-encoded byte buffer (16-bit PCM, mono, 16kHz)
samples_wav: Vec<Vec<u8>>,
// raw i16 samples kept around so we can also dump .wav keepsakes to disk
samples_pcm: Vec<Vec<i16>>,
recorder: Option<pv_recorder::PvRecorder>,
}
struct State {
session: Mutex<Option<Session>>,
}
static STATE: OnceCell<State> = OnceCell::new();
fn state() -> &'static State {
STATE.get_or_init(|| State {
session: Mutex::new(None),
})
}
/// Returns `<APP_CONFIG_DIR>/wake_words/`, creating it on first call.
pub fn models_dir() -> Result<PathBuf, String> {
let cfg = APP_CONFIG_DIR
.get()
.ok_or_else(|| "APP_CONFIG_DIR not initialised".to_string())?;
let dir = cfg.join(WAKE_WORDS_DIR_NAME);
if !dir.exists() {
fs::create_dir_all(&dir).map_err(|e| format!("create wake_words dir: {}", e))?;
}
Ok(dir)
}
pub fn status() -> TrainerStatus {
let s = state().session.lock();
if let Some(sess) = s.as_ref() {
TrainerStatus {
recording: true,
session_name: Some(sess.name.clone()),
target_samples: sess.target_samples,
collected: sess.samples_wav.len(),
}
} else {
TrainerStatus {
recording: false,
session_name: None,
target_samples: DEFAULT_TARGET_SAMPLES,
collected: 0,
}
}
}
/// Open the mic and enter a new training session. Returns Err if a session
/// already exists or the device is unavailable (e.g. jarvis-app holds it).
pub fn start(name: &str, target_samples: u8) -> Result<(), String> {
let clean_name = sanitize(name);
if clean_name.is_empty() {
return Err("Имя пустое или содержит только запрещённые символы".to_string());
}
let target = target_samples.clamp(MIN_TARGET_SAMPLES, MAX_TARGET_SAMPLES);
let mut slot = state().session.lock();
if slot.is_some() {
return Err("Сессия записи уже активна".to_string());
}
let recorder = pv_recorder::PvRecorderBuilder::new(FRAME_LEN)
.device_index(-1) // default device
.init()
.map_err(|e| {
format!(
"Не удалось открыть микрофон ({:?}). Сначала остановите J.A.R.V.I.S.",
e
)
})?;
recorder
.start()
.map_err(|e| format!("Микрофон отказал в записи: {}", e))?;
*slot = Some(Session {
name: clean_name,
target_samples: target,
samples_wav: Vec::with_capacity(target as usize),
samples_pcm: Vec::with_capacity(target as usize),
recorder: Some(recorder),
});
Ok(())
}
/// Capture one fixed-duration sample. Blocks for `seconds` (clamped to 0.5..5.0).
pub fn record_sample(seconds: f32) -> Result<usize, String> {
let dur = seconds.clamp(0.5, 5.0);
let mut slot = state().session.lock();
let sess = slot
.as_mut()
.ok_or_else(|| "Сессия не запущена".to_string())?;
if sess.samples_wav.len() >= sess.target_samples as usize {
return Err("Все сэмплы уже записаны — нажмите 'Сохранить'".to_string());
}
let recorder = sess
.recorder
.as_ref()
.ok_or_else(|| "Микрофон не инициализирован".to_string())?;
let total_samples = (SAMPLE_RATE as f32 * dur) as usize;
let mut pcm: Vec<i16> = Vec::with_capacity(total_samples + FRAME_LEN as usize);
let deadline = Instant::now() + Duration::from_millis((dur * 1000.0) as u64 + 50);
while pcm.len() < total_samples {
if Instant::now() > deadline {
// hard cap — shouldn't normally trigger but guards against driver hang
break;
}
match recorder.read() {
Ok(frame) => pcm.extend_from_slice(&frame),
Err(e) => return Err(format!("Чтение микрофона: {}", e)),
}
}
pcm.truncate(total_samples);
let wav = encode_wav(&pcm)
.map_err(|e| format!("Не удалось закодировать WAV: {}", e))?;
sess.samples_wav.push(wav);
sess.samples_pcm.push(pcm);
Ok(sess.samples_wav.len())
}
/// Train the model + persist .rpw and keepsake WAV files. Returns the path
/// to the new `.rpw`. Session is cleared on success regardless.
pub fn finish(threshold: f32) -> Result<PathBuf, String> {
let mut slot = state().session.lock();
let sess = slot
.take()
.ok_or_else(|| "Сессия не запущена".to_string())?;
if let Some(rec) = sess.recorder {
let _ = rec.stop();
}
if sess.samples_wav.len() < MIN_TARGET_SAMPLES as usize {
return Err(format!(
"Нужно минимум {} сэмплов, у тебя {}",
MIN_TARGET_SAMPLES,
sess.samples_wav.len()
));
}
let dir = models_dir()?;
let model_path = dir.join(format!("{}.rpw", sess.name));
let samples_dir = dir.join(&sess.name);
if !samples_dir.exists() {
fs::create_dir_all(&samples_dir)
.map_err(|e| format!("создать папку сэмплов: {}", e))?;
}
// dump keepsake WAVs so the user (or a future retrain) can reuse them
for (i, pcm) in sess.samples_pcm.iter().enumerate() {
let path = samples_dir.join(format!("sample_{:02}.wav", i + 1));
if let Ok(bytes) = encode_wav(pcm) {
let _ = fs::write(path, bytes);
}
}
// build the in-memory sample map rustpotter wants (file_name -> wav bytes)
let mut samples_map: HashMap<String, Vec<u8>> = HashMap::with_capacity(sess.samples_wav.len());
for (i, wav) in sess.samples_wav.into_iter().enumerate() {
samples_map.insert(format!("sample_{:02}.wav", i + 1), wav);
}
let threshold = threshold.clamp(0.3, 0.9);
let model = WakewordRef::new_from_sample_buffers(
sess.name.clone(),
Some(threshold),
None,
samples_map,
MFCC_SIZE,
)
.map_err(|e| format!("rustpotter обучение: {}", e))?;
let path_str = model_path
.to_str()
.ok_or_else(|| "Путь содержит не-UTF8".to_string())?;
model
.save_to_file(path_str)
.map_err(|e| format!("сохранение модели: {}", e))?;
info!(
"Wake-word model trained and saved: {} (threshold {:.2})",
model_path.display(),
threshold
);
Ok(model_path)
}
/// Abort the current session, releasing the mic. Always succeeds.
pub fn cancel() {
let mut slot = state().session.lock();
if let Some(sess) = slot.take() {
if let Some(rec) = sess.recorder {
let _ = rec.stop();
}
}
}
/// List all `.rpw` models the user has trained so far.
pub fn list_models() -> Vec<TrainedModelInfo> {
let dir = match models_dir() {
Ok(d) => d,
Err(_) => return Vec::new(),
};
let read = match fs::read_dir(&dir) {
Ok(r) => r,
Err(_) => return Vec::new(),
};
let mut out = Vec::new();
for entry in read.flatten() {
let path = entry.path();
if !path.is_file() {
continue;
}
if path.extension().and_then(|s| s.to_str()) != Some("rpw") {
continue;
}
let name = path
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_default();
let meta = match path.metadata() {
Ok(m) => m,
Err(_) => continue,
};
let modified = meta
.modified()
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map(|d| d.as_secs() as i64)
.unwrap_or(0);
out.push(TrainedModelInfo {
name,
path: path.display().to_string(),
size_bytes: meta.len(),
modified,
});
}
out.sort_by(|a, b| a.name.cmp(&b.name));
out
}
/// Delete a previously trained model + its keepsake samples.
pub fn delete_model(name: &str) -> Result<(), String> {
let clean = sanitize(name);
if clean.is_empty() {
return Err("имя пустое".to_string());
}
let dir = models_dir()?;
let rpw = dir.join(format!("{}.rpw", clean));
let samples = dir.join(&clean);
if rpw.is_file() {
fs::remove_file(&rpw).map_err(|e| format!("rm rpw: {}", e))?;
}
if samples.is_dir() {
fs::remove_dir_all(&samples).map_err(|e| format!("rm samples: {}", e))?;
}
Ok(())
}
fn sanitize(raw: &str) -> String {
// strip path separators + control chars; keep ascii alnum, underscore, dash, dot
raw.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.' {
c
} else if c.is_alphabetic() {
c
} else {
'_'
}
})
.collect::<String>()
.trim_matches('_')
.trim_matches('.')
.to_string()
}
fn encode_wav(pcm: &[i16]) -> Result<Vec<u8>, String> {
use std::io::{Cursor, Write};
let mut buf: Vec<u8> = Vec::with_capacity(44 + pcm.len() * 2);
let mut cursor = Cursor::new(&mut buf);
let byte_rate: u32 = SAMPLE_RATE * 2; // mono * 16-bit
let data_size: u32 = (pcm.len() * 2) as u32;
let chunk_size: u32 = 36 + data_size;
cursor.write_all(b"RIFF").map_err(|e| e.to_string())?;
cursor.write_all(&chunk_size.to_le_bytes()).map_err(|e| e.to_string())?;
cursor.write_all(b"WAVE").map_err(|e| e.to_string())?;
cursor.write_all(b"fmt ").map_err(|e| e.to_string())?;
cursor.write_all(&16u32.to_le_bytes()).map_err(|e| e.to_string())?; // PCM subchunk size
cursor.write_all(&1u16.to_le_bytes()).map_err(|e| e.to_string())?; // PCM format
cursor.write_all(&1u16.to_le_bytes()).map_err(|e| e.to_string())?; // channels
cursor.write_all(&SAMPLE_RATE.to_le_bytes()).map_err(|e| e.to_string())?;
cursor.write_all(&byte_rate.to_le_bytes()).map_err(|e| e.to_string())?;
cursor.write_all(&2u16.to_le_bytes()).map_err(|e| e.to_string())?; // block align
cursor.write_all(&16u16.to_le_bytes()).map_err(|e| e.to_string())?; // bits per sample
cursor.write_all(b"data").map_err(|e| e.to_string())?;
cursor.write_all(&data_size.to_le_bytes()).map_err(|e| e.to_string())?;
for s in pcm {
cursor.write_all(&s.to_le_bytes()).map_err(|e| e.to_string())?;
}
Ok(buf)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sanitize_strips_path_chars() {
// leading dots/underscores are also trimmed for safety
assert_eq!(sanitize("../etc/passwd"), "_etc_passwd");
assert_eq!(sanitize("my-wake"), "my-wake");
assert_eq!(sanitize("hello world"), "hello_world");
assert_eq!(sanitize("Привет.Мир"), "Привет.Мир");
assert_eq!(sanitize("___"), "");
// ensure result never starts with a dot — guards against hidden files
// and against traversal: even "..\..\foo" can't escape its directory
assert!(!sanitize("..\\..\\foo").starts_with('.'));
}
#[test]
fn wav_encoding_round_trip_size_check() {
let pcm: Vec<i16> = (0..1600).map(|i| (i % 100) as i16).collect();
let wav = encode_wav(&pcm).expect("encode wav");
// 44-byte RIFF header + 2 bytes per sample
assert_eq!(wav.len(), 44 + pcm.len() * 2);
assert_eq!(&wav[0..4], b"RIFF");
assert_eq!(&wav[8..12], b"WAVE");
assert_eq!(&wav[12..16], b"fmt ");
assert_eq!(&wav[36..40], b"data");
// sample rate at offset 24, little-endian
let sr = u32::from_le_bytes([wav[24], wav[25], wav[26], wav[27]]);
assert_eq!(sr, SAMPLE_RATE);
}
#[test]
fn idle_status_when_no_session() {
// ensure STATE init doesn't panic and reports recording=false
let s = status();
assert!(!s.recording);
assert_eq!(s.collected, 0);
}
#[test]
fn cancel_when_idle_is_a_noop() {
cancel();
// no panic, status remains idle
let s = status();
assert!(!s.recording);
}
#[test]
fn record_sample_without_session_errors() {
let r = record_sample(1.0);
assert!(r.is_err(), "must error when no session is active");
}
}