From f01154f601fd3e3609258dfd422d55c3e5278532 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 13 May 2025 05:33:18 -0400 Subject: [PATCH] persist fal keys to disk --- .../desktop/artcraft/src/commands/fal/mod.rs | 1 + .../commands/fal/set_fal_api_key_command.rs | 59 ++++++++++++++++ crates/desktop/artcraft/src/commands/mod.rs | 1 + crates/desktop/artcraft/src/lib.rs | 9 ++- .../src/state/data_dir/app_credentials_dir.rs | 4 ++ .../src/state/fal/fal_credential_manager.rs | 69 ++++++++++++++++++- crates/desktop/artcraft/src/state/fal/mod.rs | 3 +- .../state/fal/read_fal_api_key_from_disk.rs | 19 +++++ 8 files changed, 161 insertions(+), 4 deletions(-) create mode 100644 crates/desktop/artcraft/src/commands/fal/mod.rs create mode 100644 crates/desktop/artcraft/src/commands/fal/set_fal_api_key_command.rs create mode 100644 crates/desktop/artcraft/src/state/fal/read_fal_api_key_from_disk.rs diff --git a/crates/desktop/artcraft/src/commands/fal/mod.rs b/crates/desktop/artcraft/src/commands/fal/mod.rs new file mode 100644 index 0000000000..af375b9758 --- /dev/null +++ b/crates/desktop/artcraft/src/commands/fal/mod.rs @@ -0,0 +1 @@ +pub mod set_fal_api_key_command; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/commands/fal/set_fal_api_key_command.rs b/crates/desktop/artcraft/src/commands/fal/set_fal_api_key_command.rs new file mode 100644 index 0000000000..976a774307 --- /dev/null +++ b/crates/desktop/artcraft/src/commands/fal/set_fal_api_key_command.rs @@ -0,0 +1,59 @@ +use crate::commands::command_response_wrapper::CommandResult; +use crate::state::app_preferences::app_preferences_manager::AppPreferencesManager; +use crate::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory; +use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::fal::fal_credential_manager::FalCredentialManager; +use anyhow::anyhow; +use errors::AnyhowResult; +use fal_client::creds::fal_api_key::FalApiKey; +use log::{error, info}; +use openai_sora_client::sora_error::SoraError::AnyhowError; +use serde_derive::{Deserialize, Serialize}; +use tauri::{AppHandle, State}; + +/// For now, we'll only update a single preference at a time. +#[derive(Deserialize)] +pub struct SetFalApiKeyRequest{ + pub key: Option, +} + +#[tauri::command] +pub async fn set_fal_api_key_command( + request: SetFalApiKeyRequest, + creds_manager: State<'_, FalCredentialManager>, +) -> CommandResult<(), (), ()> { + info!("update_app_preferences_command called"); + + set_key(request, &creds_manager) + .await + .map_err(|err| { + error!("Error getting app preferences: {:?}", err); + "error setting key" + })?; + + Ok(().into()) +} + +async fn set_key( + request: SetFalApiKeyRequest, + creds_manager: &FalCredentialManager, +) -> AnyhowResult<()> { + + let key = request.key + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty()) + .map(|key| FalApiKey::from_str(&key)); + + match key { + None => { + creds_manager.clear_key()?; + creds_manager.purge_api_key_from_disk()?; + } + Some(key) => { + creds_manager.set_key(&key)?; + creds_manager.persist_to_disk()?; + } + } + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/commands/mod.rs b/crates/desktop/artcraft/src/commands/mod.rs index 549119c5dc..60f1daaac4 100644 --- a/crates/desktop/artcraft/src/commands/mod.rs +++ b/crates/desktop/artcraft/src/commands/mod.rs @@ -1,5 +1,6 @@ pub mod app_preferences; pub mod command_response_wrapper; +pub mod fal; pub mod flip_image; pub mod platform_info_command; pub mod sora; diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 14d7339702..0b3664f650 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -9,6 +9,7 @@ use tauri::Manager; use crate::commands::app_preferences::get_app_preferences_command::get_app_preferences_command; use crate::commands::app_preferences::update_app_preference_command::update_app_preferences_command; +use crate::commands::fal::set_fal_api_key_command::set_fal_api_key_command; use crate::commands::flip_image::flip_image; use crate::commands::platform_info_command::platform_info_command; use crate::commands::sora::check_sora_session_command::check_sora_session_command; @@ -18,6 +19,7 @@ use crate::commands::sora::sora_image_remix_command::sora_image_remix_command; use crate::commands::sora::sora_logout_command::sora_logout_command; use crate::state::app_preferences::app_preferences_manager::load_app_preferences_or_default; use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::fal::fal_credential_manager::FalCredentialManager; use crate::state::main_window_position::MainWindowPosition; use crate::state::main_window_size::MainWindowSize; use crate::state::sora::sora_credential_manager::SoraCredentialManager; @@ -48,9 +50,12 @@ pub fn run() { let storyteller_creds_manager = StorytellerCredentialManager::initialize_from_disk_infallible(&app_data_root); let storyteller_creds_manager_2 = storyteller_creds_manager.clone(); - println!("Attempting to read existing credentials..."); + println!("Attempting to read existing sora credentials..."); let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root); let sora_creds_manager_2 = sora_creds_manager.clone(); + + println!("Attempting to read existing fal credentials..."); + let fal_creds_manager = FalCredentialManager::initialize_from_disk_infallible(&app_data_root); // Other state let sora_task_queue = SoraTaskQueue::new(); @@ -134,6 +139,7 @@ pub fn run() { .manage(app_data_root) .manage(app_preferences) .manage(sora_creds_manager) + .manage(fal_creds_manager) .manage(sora_task_queue) .invoke_handler(tauri::generate_handler![ check_sora_session_command, @@ -141,6 +147,7 @@ pub fn run() { get_app_preferences_command, open_sora_login_command, platform_info_command, + set_fal_api_key_command, sora_image_generation_command, sora_image_remix_command, sora_logout_command, diff --git a/crates/desktop/artcraft/src/state/data_dir/app_credentials_dir.rs b/crates/desktop/artcraft/src/state/data_dir/app_credentials_dir.rs index cb6f51110d..2ec84cba9e 100644 --- a/crates/desktop/artcraft/src/state/data_dir/app_credentials_dir.rs +++ b/crates/desktop/artcraft/src/state/data_dir/app_credentials_dir.rs @@ -40,4 +40,8 @@ impl AppCredentialsDir { pub fn get_storyteller_session_cookie_file_path(&self) -> PathBuf { self.path.join("artcraft_session.txt") } + + pub fn get_fal_api_key_file_path(&self) -> PathBuf { + self.path.join("fal_api_key.txt") + } } diff --git a/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs b/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs index 93b572c193..cf43b1ca11 100644 --- a/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs +++ b/crates/desktop/artcraft/src/state/fal/fal_credential_manager.rs @@ -1,16 +1,43 @@ +use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::fal::read_fal_api_key_from_disk::read_fal_api_key_from_disk; +use anyhow::anyhow; use errors::AnyhowResult; use fal_client::creds::fal_api_key::FalApiKey; use memory_store::clone_slot::CloneSlot; +use std::fs::OpenOptions; +use std::io::Write; +use std::path::Path; +use std::{fs, io}; #[derive(Clone)] pub struct FalCredentialManager { - key: CloneSlot + key: CloneSlot, + app_data_root: AppDataRoot, } impl FalCredentialManager { - pub fn new() -> Self { + pub fn new(app_data_root: &AppDataRoot) -> Self { Self { key: CloneSlot::empty(), + app_data_root: app_data_root.clone(), + } + } + + pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { + let mut holder = CloneSlot::empty(); + + match read_fal_api_key_from_disk(app_data_root) { + Err(err) => { + log::warn!("Failed to read FAL API key from disk: {:?}", err); + } + Ok(api_key) => { + holder.set_clone(&api_key).expect("Failed to set FAL API key"); + } + } + + Self { + key: CloneSlot::empty(), + app_data_root: app_data_root.clone(), } } @@ -21,4 +48,42 @@ impl FalCredentialManager { pub fn get_key(&self) -> AnyhowResult> { self.key.get_clone() } + + pub fn clear_key(&self) -> AnyhowResult<()> { + self.key.clear() + } + + pub fn purge_api_key_from_disk(&self) -> Result<(), io::Error> { + let filename = self.app_data_root.credentials_dir().get_fal_api_key_file_path(); + if filename.exists() && filename.is_file() { + fs::remove_file(filename)?; + } + Ok(()) + } + + pub fn persist_to_disk(&self) -> AnyhowResult<()> { + let key = match self.get_key() { + Ok(Some(key)) => key, + Ok(None) => { + self.purge_api_key_from_disk()?; + return Ok(()); + }, + Err(err) => { + return Err(anyhow!("Failed to get FAL key: {:?}", err)); + } + }; + + let filename = self.app_data_root.credentials_dir().get_fal_api_key_file_path(); + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(filename)?; + + file.write_all(key.0.as_bytes())?; + file.flush()?; + + Ok(()) + } } diff --git a/crates/desktop/artcraft/src/state/fal/mod.rs b/crates/desktop/artcraft/src/state/fal/mod.rs index cceec1a4ce..fd55133c18 100644 --- a/crates/desktop/artcraft/src/state/fal/mod.rs +++ b/crates/desktop/artcraft/src/state/fal/mod.rs @@ -1 +1,2 @@ -pub mod fal_credential_manager; \ No newline at end of file +pub mod fal_credential_manager; +pub mod read_fal_api_key_from_disk; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/state/fal/read_fal_api_key_from_disk.rs b/crates/desktop/artcraft/src/state/fal/read_fal_api_key_from_disk.rs new file mode 100644 index 0000000000..e19ede95c0 --- /dev/null +++ b/crates/desktop/artcraft/src/state/fal/read_fal_api_key_from_disk.rs @@ -0,0 +1,19 @@ +use crate::state::data_dir::app_data_root::AppDataRoot; +use anyhow::anyhow; +use errors::AnyhowResult; +use fal_client::creds::fal_api_key::FalApiKey; +use std::fs::read_to_string; + +pub fn read_fal_api_key_from_disk(app_data_root: &AppDataRoot) -> AnyhowResult { + let api_key_file = app_data_root.credentials_dir().get_fal_api_key_file_path(); + + if !api_key_file.exists() { + return Err(anyhow!("API key file does not exist: {:?}", api_key_file)); + } + + let value = read_to_string(api_key_file)? + .trim() + .to_string(); + + Ok(FalApiKey::new(value)) +}