From fe9dfbb63a37d5a1ffd8db94ea835319039aa3a2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 21 Apr 2025 22:52:18 -0400 Subject: [PATCH] better credential behavior --- .../commands/sora/sora_image_remix_command.rs | 60 +++++++++++--- crates/desktop/tauri-artcraft/src/lib.rs | 16 ++-- .../tauri-artcraft/src/state/sora/mod.rs | 3 +- .../src/state/sora/sora_credential_holder.rs | 17 +++- .../src/state/sora/sora_credential_manager.rs | 83 +++++++++++++++++++ .../src/threads/sora_session_login_thread.rs | 3 +- 6 files changed, 156 insertions(+), 26 deletions(-) create mode 100644 crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs index 6f70c05fda..11c8d650de 100644 --- a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs +++ b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs @@ -1,8 +1,9 @@ -use base64::Engine; -use base64::prelude::BASE64_STANDARD; use crate::state::app_dir::AppDataRoot; use crate::state::sora::read_sora_credentials_from_disk::read_sora_credentials_from_disk; use crate::state::sora::sora_credential_holder::SoraCredentialHolder; +use crate::state::sora::sora_credential_manager::SoraCredentialManager; +use base64::prelude::BASE64_STANDARD; +use base64::Engine; use errors::AnyhowResult; use image::codecs::png::{CompressionType, FilterType, PngEncoder}; use image::{DynamicImage, ImageReader}; @@ -10,6 +11,7 @@ use log::{error, info}; use openai_sora_client::credentials::SoraCredentials; use openai_sora_client::image_gen::common::{ImageSize, NumImages}; use openai_sora_client::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest}; +use openai_sora_client::sentinel_refresh::refresh_sentinel; use openai_sora_client::upload::upload_media_from_bytes::sora_media_upload_from_bytes; use openai_sora_client::upload::upload_media_from_file::SoraMediaUploadRequest; use serde_derive::Deserialize; @@ -40,11 +42,11 @@ pub async fn sora_image_remix_command( _app: AppHandle, request: SoraImageRemixCommand, app_data_root: State<'_, AppDataRoot>, - sora_creds_holder: State<'_, SoraCredentialHolder>, + sora_creds_manager: State<'_, SoraCredentialManager>, ) -> Result { info!("image_generation_command called; processing image..."); - generate_image(request, &app_data_root, &sora_creds_holder) + generate_image(request, &app_data_root, &sora_creds_manager) .await .map_err(|err| { error!("error: {:?}", err); @@ -57,16 +59,31 @@ pub async fn sora_image_remix_command( pub async fn generate_image( request: SoraImageRemixCommand, app_data_root: &AppDataRoot, - sora_creds_holder: &SoraCredentialHolder, + sora_creds_manager: &SoraCredentialManager, ) -> AnyhowResult<()> { - let sora_credentials = read_sora_credentials_from_disk(app_data_root) + // TODO(bt,2025-04-21): Read from in-memory cache instead, but allow for desktop replacement. + let mut sora_creds = read_sora_credentials_from_disk(app_data_root) .map_err(|err| { error!("Failed to read Sora credentials from disk: {:?}", err); err })?; - //let sora_credentials = sora_creds_holder.get_credentials()?; + let no_sentinel = !app_data_root.get_sora_sentinel_file_path().is_file() + || sora_creds.sentinel.is_none(); + + if no_sentinel { + info!("Refreshing credentials..."); + + sora_creds = sora_creds_manager.call_sentinel_refresh() + .await + .map_err(|err| { + error!("Failed to refresh: {:?}", err); + err + })?; + + // TODO: Write to disk. + } let sora_media_tokens = vec![]; @@ -74,13 +91,36 @@ pub async fn generate_image( // Note: This is incredibly inefficient. We should keep a local cache. // Also, if they've already been uploaded to OpenAI, we shouldn't continue to re-upload. - let response = sora_image_gen_remix(SoraImageGenRemixRequest { + let mut response = sora_image_gen_remix(SoraImageGenRemixRequest { prompt: request.prompt.to_string(), num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &sora_credentials, - }).await + credentials: &sora_creds, + }).await; + + if let Err(err) = &response { + error!("Error in generating image: {:?}", err); + + sora_creds = sora_creds_manager.call_sentinel_refresh() + .await + .map_err(|err| { + error!("Failed to refresh: {:?}", err); + err + })?; + + info!("Retrying request..."); + + response = sora_image_gen_remix(SoraImageGenRemixRequest { + prompt: request.prompt.to_string(), + num_images: NumImages::One, + image_size: ImageSize::Square, + sora_media_tokens: sora_media_tokens.clone(), + credentials: &sora_creds, + }).await; + } + + let response = response .map_err(|err| { error!("Failed to call Sora image generation: {:?}", err); err diff --git a/crates/desktop/tauri-artcraft/src/lib.rs b/crates/desktop/tauri-artcraft/src/lib.rs index a191fb8fa8..578d47679a 100644 --- a/crates/desktop/tauri-artcraft/src/lib.rs +++ b/crates/desktop/tauri-artcraft/src/lib.rs @@ -11,6 +11,7 @@ use crate::commands::sora::sora_image_generation_command::sora_image_generation_ use crate::commands::sora::sora_image_remix_command::sora_image_remix_command; use crate::state::app_config::AppConfig; use crate::state::sora::sora_credential_holder::SoraCredentialHolder; +use crate::state::sora::sora_credential_manager::SoraCredentialManager; use crate::threads::sora_session_login_thread::sora_session_login_thread; use tauri_plugin_log::Target; @@ -29,16 +30,9 @@ pub fn run() { let app_data_root = config.app_data_root.clone(); let app_data_root2 = config.app_data_root.clone(); - let sora_creds_holder = SoraCredentialHolder::new(); - let sora_creds_holder2 = sora_creds_holder.clone(); - println!("Attempting to read existing credentials..."); - - if let Ok(creds) = read_sora_credentials_from_disk(&app_data_root) { - sora_creds_holder.set_credentials(creds).expect("there shouldn't be locking errors"); - } else { - println!("No credentials found on disk."); - } + let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root); + let sora_creds_manager2 = sora_creds_manager.clone(); println!("Initializing backend runtime..."); @@ -65,13 +59,13 @@ pub fn run() { //} let app = app.handle().clone(); - tauri::async_runtime::spawn(sora_session_login_thread(app, app_data_root2, sora_creds_holder2)); + tauri::async_runtime::spawn(sora_session_login_thread(app, app_data_root2, sora_creds_manager2)); Ok(()) }) .manage(config) .manage(app_data_root) - .manage(sora_creds_holder) + .manage(sora_creds_manager) .invoke_handler(tauri::generate_handler![ flip_image, open_sora_login_command, diff --git a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs index 0ce3ff2d4d..0726ea796b 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs @@ -1,2 +1,3 @@ pub mod read_sora_credentials_from_disk; -pub mod sora_credential_holder; \ No newline at end of file +pub mod sora_credential_holder; +pub mod sora_credential_manager; \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs index 649efe1093..d86d4f156c 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs @@ -1,4 +1,5 @@ -use std::sync::{Arc, RwLock}; +use std::ops::Deref; +use std::sync::{Arc, RwLock, RwLockReadGuard}; use anyhow::anyhow; use errors::AnyhowResult; use openai_sora_client::credentials::SoraCredentials; @@ -16,11 +17,11 @@ impl SoraCredentialHolder { } } - pub fn set_credentials(&self, credentials: SoraCredentials) -> AnyhowResult<()> { + pub fn set_credentials(&self, credentials: &SoraCredentials) -> AnyhowResult<()> { match self.credentials.write() { Err(err) => Err(anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut creds) => { - *creds = Some(credentials); + *creds = Some(credentials.clone()); Ok(()) } } @@ -44,4 +45,14 @@ impl SoraCredentialHolder { } } } + + pub fn get_credentials_required(&self) -> AnyhowResult { + match self.credentials.read() { + Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(creds) => match creds.deref() { + None => Err(anyhow!("Credentials not set")), + Some(creds) => Ok(creds.clone()), + } + } + } } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs new file mode 100644 index 0000000000..fcb63b8109 --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs @@ -0,0 +1,83 @@ +use crate::state::app_dir::AppDataRoot; +use crate::state::sora::read_sora_credentials_from_disk::read_sora_credentials_from_disk; +use crate::state::sora::sora_credential_holder::SoraCredentialHolder; +use errors::AnyhowResult; +use log::{error, warn}; +use openai_sora_client::credentials::SoraCredentials; +use openai_sora_client::sentinel_refresh::refresh_sentinel; +use std::fs::OpenOptions; +use std::io::Write; + +#[derive(Clone)] +pub struct SoraCredentialManager { + holder: SoraCredentialHolder, + app_data_root: AppDataRoot, +} + +impl SoraCredentialManager { + + pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { + let holder = SoraCredentialHolder::new(); + + match read_sora_credentials_from_disk(app_data_root) { + Err(err) => warn!("Failed to read credentials from disk: {:?}", err), + Ok(creds) => { + holder.set_credentials(&creds).expect("Failed to set credentials"); + } + } + + Self { + holder, + app_data_root: app_data_root.clone(), + } + } + + pub fn set_credentials(&self, creds: &SoraCredentials) -> AnyhowResult<()> { + self.holder.set_credentials(creds) + } + + pub fn clear_credentials(&self) -> AnyhowResult<()> { + self.holder.clear_credentials() + } + + pub fn get_credentials(&self) -> AnyhowResult> { + self.holder.get_credentials() + } + + pub fn get_credentials_required(&self) -> AnyhowResult { + self.holder.get_credentials_required() + } + + pub fn reset_from_disk(&self) -> AnyhowResult { + let creds = read_sora_credentials_from_disk(&self.app_data_root)?; + self.holder.set_credentials(&creds)?; + Ok(creds) + } + + pub async fn call_sentinel_refresh(&self) -> AnyhowResult { + // NB(bt,2025-04-21): Technically we don't need credentials to get a sentinel. + let mut creds = self.holder.get_credentials_required()?; + + let sentinel = refresh_sentinel() + .await + .map_err(|err| { + error!("Failed to refresh: {:?}", err); + err + })?; + + creds.sentinel = Some(sentinel.clone()); + + self.holder.set_credentials(&creds)?; + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(self.app_data_root.get_sora_sentinel_file_path())?; + + file.write_all(sentinel.as_bytes())?; + file.flush()?; + + Ok(creds) + } +} diff --git a/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs b/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs index 2af700f397..c8bb5befb1 100644 --- a/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs +++ b/crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs @@ -8,6 +8,7 @@ use tauri::{AppHandle, Manager, Webview}; use errors::AnyhowResult; use crate::state::app_dir::AppDataRoot; use crate::state::sora::sora_credential_holder::SoraCredentialHolder; +use crate::state::sora::sora_credential_manager::SoraCredentialManager; use crate::utils::sora_webview_cookies::get_all_sora_cookies_as_string; pub const LOGIN_WINDOW_NAME: &str = "login_window"; @@ -27,7 +28,7 @@ pub static SORA_ROOT_URL: Lazy = Lazy::new(|| { pub async fn sora_session_login_thread( app: AppHandle, app_data_root: AppDataRoot, - sora_creds_holder: SoraCredentialHolder + sora_creds_manager: SoraCredentialManager ) -> ! { loop { for (window_name, webview) in app.webviews() {