//! Tauri-specific extensions for yaak-models. //! //! This module provides the Tauri plugin initialization and extension traits //! that allow accessing QueryManager and BlobManager from Tauri's Manager types. use chrono::Utc; use log::error; use std::time::Duration; use tauri::plugin::TauriPlugin; use tauri::{Emitter, Manager, Runtime, State}; use tauri_plugin_dialog::{DialogExt, MessageDialogKind}; use yaak_models::blob_manager::BlobManager; use yaak_models::db_context::DbContext; use yaak_models::error::Result; use yaak_models::models::{AnyModel, GraphQlIntrospection, GrpcEvent, Settings, WebsocketEvent}; use yaak_models::query_manager::QueryManager; use yaak_models::util::UpdateSource; const MODEL_CHANGES_RETENTION_HOURS: i64 = 1; const MODEL_CHANGES_POLL_INTERVAL_MS: u64 = 1000; const MODEL_CHANGES_POLL_BATCH_SIZE: usize = 200; struct ModelChangeCursor { created_at: String, id: i64, } impl ModelChangeCursor { fn from_launch_time() -> Self { Self { created_at: Utc::now().naive_utc().format("%Y-%m-%d %H:%M:%S%.3f").to_string(), id: 0, } } } fn drain_model_changes_batch( query_manager: &QueryManager, app_handle: &tauri::AppHandle, cursor: &mut ModelChangeCursor, ) -> bool { let changes = match query_manager.connect().list_model_changes_since( &cursor.created_at, cursor.id, MODEL_CHANGES_POLL_BATCH_SIZE, ) { Ok(changes) => changes, Err(err) => { error!("Failed to poll model_changes rows: {err:?}"); return false; } }; if changes.is_empty() { return false; } let fetched_count = changes.len(); for change in changes { cursor.created_at = change.created_at; cursor.id = change.id; // Local window-originated writes are forwarded immediately from the // in-memory model event channel. if matches!(change.payload.update_source, UpdateSource::Window { .. }) { continue; } if let Err(err) = app_handle.emit("model_write", change.payload) { error!("Failed to emit model_write event: {err:?}"); } } fetched_count == MODEL_CHANGES_POLL_BATCH_SIZE } async fn run_model_change_poller( query_manager: QueryManager, app_handle: tauri::AppHandle, mut cursor: ModelChangeCursor, ) { loop { while drain_model_changes_batch(&query_manager, &app_handle, &mut cursor) {} tokio::time::sleep(Duration::from_millis(MODEL_CHANGES_POLL_INTERVAL_MS)).await; } } /// Extension trait for accessing the QueryManager from Tauri Manager types. pub trait QueryManagerExt<'a, R> { fn db_manager(&'a self) -> State<'a, QueryManager>; fn db(&'a self) -> DbContext<'a>; fn with_tx(&'a self, func: F) -> Result where F: FnOnce(&DbContext) -> Result; } impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { fn db_manager(&'a self) -> State<'a, QueryManager> { self.state::() } fn db(&'a self) -> DbContext<'a> { let qm = self.state::(); qm.inner().connect() } fn with_tx(&'a self, func: F) -> Result where F: FnOnce(&DbContext) -> Result, { let qm = self.state::(); qm.inner().with_tx(func) } } /// Extension trait for accessing the BlobManager from Tauri Manager types. pub trait BlobManagerExt<'a, R> { fn blob_manager(&'a self) -> State<'a, BlobManager>; fn blobs(&'a self) -> yaak_models::blob_manager::BlobContext; } impl<'a, R: Runtime, M: Manager> BlobManagerExt<'a, R> for M { fn blob_manager(&'a self) -> State<'a, BlobManager> { self.state::() } fn blobs(&'a self) -> yaak_models::blob_manager::BlobContext { let manager = self.state::(); manager.inner().connect() } } // Commands for yaak-models use tauri::WebviewWindow; #[tauri::command] pub(crate) fn models_upsert( window: WebviewWindow, model: AnyModel, ) -> Result { use yaak_models::error::Error::GenericError; let db = window.db(); let blobs = window.blob_manager(); let source = &UpdateSource::from_window_label(window.label()); let id = match model { AnyModel::CookieJar(m) => db.upsert_cookie_jar(&m, source)?.id, AnyModel::Environment(m) => db.upsert_environment(&m, source)?.id, AnyModel::Folder(m) => db.upsert_folder(&m, source)?.id, AnyModel::GrpcRequest(m) => db.upsert_grpc_request(&m, source)?.id, AnyModel::HttpRequest(m) => db.upsert_http_request(&m, source)?.id, AnyModel::HttpResponse(m) => db.upsert_http_response(&m, source, &blobs)?.id, AnyModel::KeyValue(m) => db.upsert_key_value(&m, source)?.id, AnyModel::Plugin(m) => db.upsert_plugin(&m, source)?.id, AnyModel::Settings(m) => db.upsert_settings(&m, source)?.id, AnyModel::WebsocketRequest(m) => db.upsert_websocket_request(&m, source)?.id, AnyModel::Workspace(m) => db.upsert_workspace(&m, source)?.id, AnyModel::WorkspaceMeta(m) => db.upsert_workspace_meta(&m, source)?.id, a => return Err(GenericError(format!("Cannot upsert AnyModel {a:?})"))), }; Ok(id) } #[tauri::command] pub(crate) fn models_delete( window: WebviewWindow, model: AnyModel, ) -> Result { use yaak_models::error::Error::GenericError; let blobs = window.blob_manager(); // Use transaction for deletions because it might recurse window.with_tx(|tx| { let source = &UpdateSource::from_window_label(window.label()); let id = match model { AnyModel::CookieJar(m) => tx.delete_cookie_jar(&m, source)?.id, AnyModel::Environment(m) => tx.delete_environment(&m, source)?.id, AnyModel::Folder(m) => tx.delete_folder(&m, source)?.id, AnyModel::GrpcConnection(m) => tx.delete_grpc_connection(&m, source)?.id, AnyModel::GrpcRequest(m) => tx.delete_grpc_request(&m, source)?.id, AnyModel::HttpRequest(m) => tx.delete_http_request(&m, source)?.id, AnyModel::HttpResponse(m) => tx.delete_http_response(&m, source, &blobs)?.id, AnyModel::Plugin(m) => tx.delete_plugin(&m, source)?.id, AnyModel::WebsocketConnection(m) => tx.delete_websocket_connection(&m, source)?.id, AnyModel::WebsocketRequest(m) => tx.delete_websocket_request(&m, source)?.id, AnyModel::Workspace(m) => tx.delete_workspace(&m, source)?.id, a => return Err(GenericError(format!("Cannot delete AnyModel {a:?})"))), }; Ok(id) }) } #[tauri::command] pub(crate) fn models_duplicate( window: WebviewWindow, model: AnyModel, ) -> Result { use yaak_models::error::Error::GenericError; // Use transaction for duplications because it might recurse window.with_tx(|tx| { let source = &UpdateSource::from_window_label(window.label()); let id = match model { AnyModel::Environment(m) => tx.duplicate_environment(&m, source)?.id, AnyModel::Folder(m) => tx.duplicate_folder(&m, source)?.id, AnyModel::GrpcRequest(m) => tx.duplicate_grpc_request(&m, source)?.id, AnyModel::HttpRequest(m) => tx.duplicate_http_request(&m, source)?.id, AnyModel::WebsocketRequest(m) => tx.duplicate_websocket_request(&m, source)?.id, a => return Err(GenericError(format!("Cannot duplicate AnyModel {a:?})"))), }; Ok(id) }) } #[tauri::command] pub(crate) fn models_websocket_events( app_handle: tauri::AppHandle, connection_id: &str, ) -> Result> { Ok(app_handle.db().list_websocket_events(connection_id)?) } #[tauri::command] pub(crate) fn models_grpc_events( app_handle: tauri::AppHandle, connection_id: &str, ) -> Result> { Ok(app_handle.db().list_grpc_events(connection_id)?) } #[tauri::command] pub(crate) fn models_get_settings(app_handle: tauri::AppHandle) -> Result { Ok(app_handle.db().get_settings()) } #[tauri::command] pub(crate) fn models_get_graphql_introspection( app_handle: tauri::AppHandle, request_id: &str, ) -> Result> { Ok(app_handle.db().get_graphql_introspection(request_id)) } #[tauri::command] pub(crate) fn models_upsert_graphql_introspection( app_handle: tauri::AppHandle, request_id: &str, workspace_id: &str, content: Option, window: WebviewWindow, ) -> Result { let source = UpdateSource::from_window_label(window.label()); Ok(app_handle.db().upsert_graphql_introspection(workspace_id, request_id, content, &source)?) } #[tauri::command] pub(crate) fn models_workspace_models( window: WebviewWindow, workspace_id: Option<&str>, ) -> Result { let db = window.db(); let mut l: Vec = Vec::new(); // Add the settings l.push(db.get_settings().into()); // Add global models l.append(&mut db.list_workspaces()?.into_iter().map(Into::into).collect()); l.append(&mut db.list_key_values()?.into_iter().map(Into::into).collect()); l.append(&mut db.list_plugins()?.into_iter().map(Into::into).collect()); // Add the workspace children if let Some(wid) = workspace_id { l.append(&mut db.list_cookie_jars(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_environments_ensure_base(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_folders(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_grpc_connections(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_grpc_requests(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_http_requests(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_http_responses(wid, None)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_websocket_connections(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_websocket_requests(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_workspace_metas(wid)?.into_iter().map(Into::into).collect()); } let j = serde_json::to_string(&l)?; Ok(escape_str_for_webview(&j)) } fn escape_str_for_webview(input: &str) -> String { input .chars() .map(|c| { let code = c as u32; // ASCII if code <= 0x7F { c.to_string() // BMP characters encoded normally } else if code < 0xFFFF { format!("\\u{:04X}", code) // Beyond BMP encoded a surrogate pairs } else { let high = ((code - 0x10000) >> 10) + 0xD800; let low = ((code - 0x10000) & 0x3FF) + 0xDC00; format!("\\u{:04X}\\u{:04X}", high, low) } }) .collect() } /// Initialize database managers as a plugin (for initialization order). /// Commands are in the main invoke_handler. /// This must be registered before other plugins that depend on the database. pub fn init() -> TauriPlugin { tauri::plugin::Builder::new("yaak-models-db") .setup(|app_handle, _api| { let app_path = app_handle.path().app_data_dir().unwrap(); let db_path = app_path.join("db.sqlite"); let blob_path = app_path.join("blobs.sqlite"); let (query_manager, blob_manager, rx) = match yaak_models::init_standalone(&db_path, &blob_path) { Ok(result) => result, Err(e) => { app_handle .dialog() .message(e.to_string()) .kind(MessageDialogKind::Error) .blocking_show(); return Err(Box::from(e.to_string())); } }; let db = query_manager.connect(); if let Err(err) = db.prune_model_changes_older_than_hours(MODEL_CHANGES_RETENTION_HOURS) { error!("Failed to prune model_changes rows on startup: {err:?}"); } // Only stream writes that happen after this app launch. let cursor = ModelChangeCursor::from_launch_time(); let poll_query_manager = query_manager.clone(); app_handle.manage(query_manager); app_handle.manage(blob_manager); // Poll model_changes so all writers (including external CLI processes) update the UI. let app_handle_poll = app_handle.clone(); let query_manager = poll_query_manager; tauri::async_runtime::spawn(async move { run_model_change_poller(query_manager, app_handle_poll, cursor).await; }); // Fast path for local app writes initiated by frontend windows. This keeps the // current sync-model UX snappy, while DB polling handles external writers (CLI). let app_handle_local = app_handle.clone(); tauri::async_runtime::spawn(async move { for payload in rx { if !matches!(payload.update_source, UpdateSource::Window { .. }) { continue; } if let Err(err) = app_handle_local.emit("model_write", payload) { error!("Failed to emit local model_write event: {err:?}"); } } }); Ok(()) }) .build() }