Merge branch 'main' of github.com:cjpais/Handy

Author: cjpais

src-tauri/src/commands/mod.rs

@@ -1,5 +1,6 @@
 pub mod audio;
 pub mod models;
+pub mod transcription;
 
 use crate::utils::cancel_current_operation;
 use tauri::{AppHandle, Manager};

src-tauri/src/commands/transcription.rs

@@ -0,0 +1,32 @@
+use crate::managers::transcription::TranscriptionManager;
+use crate::settings::{get_settings, write_settings, ModelUnloadTimeout};
+use tauri::{AppHandle, State};
+
+#[tauri::command]
+pub fn set_model_unload_timeout(app: AppHandle, timeout: ModelUnloadTimeout) {
+    let mut settings = get_settings(&app);
+    settings.model_unload_timeout = timeout;
+    write_settings(&app, settings);
+}
+
+#[tauri::command]
+pub fn get_model_load_status(
+    transcription_manager: State<TranscriptionManager>,
+) -> Result<serde_json::Value, String> {
+    let is_loaded = transcription_manager.is_model_loaded();
+    let current_model = transcription_manager.get_current_model();
+
+    Ok(serde_json::json!({
+        "is_loaded": is_loaded,
+        "current_model": current_model
+    }))
+}
+
+#[tauri::command]
+pub fn unload_model_manually(
+    transcription_manager: State<TranscriptionManager>,
+) -> Result<(), String> {
+    transcription_manager
+        .unload_model()
+        .map_err(|e| format!("Failed to unload model: {}", e))
+}

src-tauri/src/lib.rs

@@ -83,20 +83,10 @@ pub fn run() {
         .manage(Mutex::new(ShortcutToggleStates::default()))
         .setup(move |app| {
             // Get the current theme to set the appropriate initial icon
-            let initial_theme = if let Some(main_window) = app.get_webview_window("main") {
-                main_window.theme().unwrap_or(tauri::Theme::Dark)
-            } else {
-                tauri::Theme::Dark
-            };
-
-            println!("Initial system theme: {:?}", initial_theme);
+            let initial_theme = tray::get_current_theme(&app.handle());
 
             // Choose the appropriate initial icon based on theme
-            let initial_icon_path = match initial_theme {
-                tauri::Theme::Dark => "resources/tray_idle.png",
-                tauri::Theme::Light => "resources/tray_idle_dark.png",
-                _ => "resources/tray_idle.png", // Default fallback
-            };
+            let initial_icon_path = tray::get_icon_path(initial_theme, tray::TrayIconState::Idle);
 
             let tray = TrayIconBuilder::new()
                 .icon(Image::from_path(app.path().resolve(
@@ -213,7 +203,10 @@ pub fn run() {
             commands::audio::get_selected_microphone,
             commands::audio::get_available_output_devices,
             commands::audio::set_selected_output_device,
-            commands::audio::get_selected_output_device
+            commands::audio::get_selected_output_device,
+            commands::transcription::set_model_unload_timeout,
+            commands::transcription::get_model_load_status,
+            commands::transcription::unload_model_manually
         ])
         .run(tauri::generate_context!())
         .expect("error while running tauri application");

src-tauri/src/managers/transcription.rs

@@ -1,9 +1,13 @@
 use crate::managers::model::ModelManager;
-use crate::settings::get_settings;
+use crate::settings::{get_settings, ModelUnloadTimeout};
 use anyhow::Result;
+use log::debug;
 use natural::phonetics::soundex;
 use serde::Serialize;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
+use std::thread;
+use std::time::{Duration, SystemTime};
 use strsim::levenshtein;
 use tauri::{App, AppHandle, Emitter, Manager};
 use whisper_rs::{
@@ -18,12 +22,16 @@ pub struct ModelStateEvent {
     pub error: Option<String>,
 }
 
+#[derive(Clone)]
 pub struct TranscriptionManager {
-    state: Mutex<Option<WhisperState>>,
-    context: Mutex<Option<WhisperContext>>,
+    state: Arc<Mutex<Option<WhisperState>>>,
+    context: Arc<Mutex<Option<WhisperContext>>>,
     model_manager: Arc<ModelManager>,
     app_handle: AppHandle,
-    current_model_id: Mutex<Option<String>>,
+    current_model_id: Arc<Mutex<Option<String>>>,
+    last_activity: Arc<AtomicU64>,
+    shutdown_signal: Arc<AtomicBool>,
+    watcher_handle: Arc<Mutex<Option<thread::JoinHandle<()>>>>,
 }
 
 fn apply_custom_words(text: &str, custom_words: &[String], threshold: f64) -> String {
@@ -139,21 +147,133 @@ impl TranscriptionManager {
         let app_handle = app.app_handle().clone();
 
         let manager = Self {
-            state: Mutex::new(None),
-            context: Mutex::new(None),
+            state: Arc::new(Mutex::new(None)),
+            context: Arc::new(Mutex::new(None)),
             model_manager,
             app_handle: app_handle.clone(),
-            current_model_id: Mutex::new(None),
+            current_model_id: Arc::new(Mutex::new(None)),
+            last_activity: Arc::new(AtomicU64::new(
+                SystemTime::now()
+                    .duration_since(SystemTime::UNIX_EPOCH)
+                    .unwrap()
+                    .as_millis() as u64,
+            )),
+            shutdown_signal: Arc::new(AtomicBool::new(false)),
+            watcher_handle: Arc::new(Mutex::new(None)),
         };
 
+        // Start the idle watcher
+        {
+            let app_handle_cloned = app_handle.clone();
+            let manager_cloned = manager.clone();
+            let shutdown_signal = manager.shutdown_signal.clone();
+            let handle = thread::spawn(move || {
+                while !shutdown_signal.load(Ordering::Relaxed) {
+                    thread::sleep(Duration::from_secs(10)); // Check every 10 seconds
+
+                    // Check shutdown signal again after sleep
+                    if shutdown_signal.load(Ordering::Relaxed) {
+                        break;
+                    }
+
+                    let settings = get_settings(&app_handle_cloned);
+                    let timeout_seconds = settings.model_unload_timeout.to_seconds();
+
+                    if let Some(limit_seconds) = timeout_seconds {
+                        // Skip polling-based unloading for immediate timeout since it's handled directly in transcribe()
+                        if settings.model_unload_timeout == ModelUnloadTimeout::Immediately {
+                            continue;
+                        }
+
+                        let last = manager_cloned.last_activity.load(Ordering::Relaxed);
+                        let now_ms = SystemTime::now()
+                            .duration_since(SystemTime::UNIX_EPOCH)
+                            .unwrap()
+                            .as_millis() as u64;
+
+                        if now_ms.saturating_sub(last) > limit_seconds * 1000 {
+                            // idle -> unload
+                            if manager_cloned.is_model_loaded() {
+                                let unload_start = std::time::Instant::now();
+                                debug!("Starting to unload model due to inactivity");
+
+                                if let Ok(()) = manager_cloned.unload_model() {
+                                    let _ = app_handle_cloned.emit(
+                                        "model-state-changed",
+                                        ModelStateEvent {
+                                            event_type: "unloaded_due_to_idle".to_string(),
+                                            model_id: None,
+                                            model_name: None,
+                                            error: None,
+                                        },
+                                    );
+                                    let unload_duration = unload_start.elapsed();
+                                    debug!(
+                                        "Model unloaded due to inactivity (took {}ms)",
+                                        unload_duration.as_millis()
+                                    );
+                                }
+                            }
+                        }
+                    }
+                }
+                debug!("Idle watcher thread shutting down gracefully");
+            });
+            *manager.watcher_handle.lock().unwrap() = Some(handle);
+        }
+
         // Try to load the default model from settings, but don't fail if no models are available
         let settings = get_settings(&app_handle);
         let _ = manager.load_model(&settings.selected_model);
 
         Ok(manager)
     }
 
+    pub fn is_model_loaded(&self) -> bool {
+        let state = self.state.lock().unwrap();
+        state.is_some()
+    }
+
+    pub fn unload_model(&self) -> Result<()> {
+        let unload_start = std::time::Instant::now();
+        debug!("Starting to unload model");
+
+        {
+            let mut state = self.state.lock().unwrap();
+            *state = None; // Dropping state frees GPU/CPU memory
+        }
+        {
+            let mut context = self.context.lock().unwrap();
+            *context = None; // Dropping context frees additional memory
+        }
+        {
+            let mut current_model = self.current_model_id.lock().unwrap();
+            *current_model = None;
+        }
+
+        // Emit unloaded event
+        let _ = self.app_handle.emit(
+            "model-state-changed",
+            ModelStateEvent {
+                event_type: "unloaded_manually".to_string(),
+                model_id: None,
+                model_name: None,
+                error: None,
+            },
+        );
+
+        let unload_duration = unload_start.elapsed();
+        debug!(
+            "Model unloaded manually (took {}ms)",
+            unload_duration.as_millis()
+        );
+        Ok(())
+    }
+
     pub fn load_model(&self, model_id: &str) -> Result<()> {
+        let load_start = std::time::Instant::now();
+        debug!("Starting to load model: {}", model_id);
+
         // Emit loading started event
         let _ = self.app_handle.emit(
             "model-state-changed",
@@ -252,7 +372,12 @@ impl TranscriptionManager {
             },
         );
 
-        println!("Successfully loaded transcription model: {}", model_id);
+        let load_duration = load_start.elapsed();
+        debug!(
+            "Successfully loaded transcription model: {} (took {}ms)",
+            model_id,
+            load_duration.as_millis()
+        );
         Ok(())
     }
 
@@ -262,6 +387,15 @@ impl TranscriptionManager {
     }
 
     pub fn transcribe(&self, audio: Vec<f32>) -> Result<String> {
+        // Update last activity timestamp
+        self.last_activity.store(
+            SystemTime::now()
+                .duration_since(SystemTime::UNIX_EPOCH)
+                .unwrap()
+                .as_millis() as u64,
+            Ordering::Relaxed,
+        );
+
         let st = std::time::Instant::now();
 
         let mut result = String::new();
@@ -272,10 +406,34 @@ impl TranscriptionManager {
             return Ok(result);
         }
 
+        // Check if model is loaded, if not try to load it
+        {
+            let state_guard = self.state.lock().unwrap();
+            if state_guard.is_none() {
+                // Model not loaded, try to load the selected model from settings
+                let settings = get_settings(&self.app_handle);
+                println!(
+                    "Model not loaded, attempting to load: {}",
+                    settings.selected_model
+                );
+
+                // Drop the guard before calling load_model to avoid deadlock
+                drop(state_guard);
+
+                // Try to load the model
+                if let Err(e) = self.load_model(&settings.selected_model) {
+                    return Err(anyhow::anyhow!(
+                        "Failed to auto-load model '{}': {}. Please check that the model is downloaded and try again.",
+                        settings.selected_model, e
+                    ));
+                }
+            }
+        }
+
         let mut state_guard = self.state.lock().unwrap();
         let state = state_guard.as_mut().ok_or_else(|| {
             anyhow::anyhow!(
-                "No model loaded. Please download and select a model from settings first."
+                "Model failed to load after auto-load attempt. Please check your model settings."
             )
         })?;
 
@@ -333,6 +491,34 @@ impl TranscriptionManager {
         };
         println!("\ntook {}ms{}", (et - st).as_millis(), translation_note);
 
+        // Check if we should immediately unload the model after transcription
+        if settings.model_unload_timeout == ModelUnloadTimeout::Immediately {
+            println!("⚡ Immediately unloading model after transcription");
+            // Drop the state guard first to avoid deadlock
+            drop(state_guard);
+            if let Err(e) = self.unload_model() {
+                eprintln!("Failed to immediately unload model: {}", e);
+            }
+        }
+
         Ok(corrected_result.trim().to_string())
     }
 }
+
+impl Drop for TranscriptionManager {
+    fn drop(&mut self) {
+        debug!("Shutting down TranscriptionManager");
+
+        // Signal the watcher thread to shutdown
+        self.shutdown_signal.store(true, Ordering::Relaxed);
+
+        // Wait for the thread to finish gracefully
+        if let Some(handle) = self.watcher_handle.lock().unwrap().take() {
+            if let Err(e) = handle.join() {
+                eprintln!("Failed to join idle watcher thread: {:?}", e);
+            } else {
+                debug!("Idle watcher thread joined successfully");
+            }
+        }
+    }
+}