From 1553465850522cd6cac3f2ec5c99318eae7fc796 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 24 Apr 2025 10:18:56 -0400 Subject: [PATCH] reset in-memory store from login thread --- .../sora/sora_image_generation_command.rs | 1 - .../commands/sora/sora_image_remix_command.rs | 3 +- crates/desktop/tauri-artcraft/src/lib.rs | 2 +- .../tauri-artcraft/src/state/sora/mod.rs | 1 + .../sora/read_sora_credentials_from_disk.rs | 52 +---------- ...sora_credentials_from_disk_with_upgrade.rs | 92 +++++++++++++++++++ .../src/state/sora/sora_credential_manager.rs | 8 +- .../src/threads/sora_session_login_thread.rs | 38 +++++++- 8 files changed, 137 insertions(+), 60 deletions(-) create mode 100644 crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk_with_upgrade.rs diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs index f90ee163a4..5dfd603fa1 100644 --- a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs +++ b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs @@ -45,7 +45,6 @@ pub async fn generate_image( ) -> AnyhowResult<()> { let creds = read_sora_credentials_from_disk(app_data_root) - .await .map_err(|err| { error!("Failed to read Sora credentials from disk: {:?}", err); err 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 cece13ad45..fb47e76f8f 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 @@ -50,6 +50,8 @@ pub async fn sora_image_remix_command( ) -> Result { info!("image_generation_command called; processing image..."); + // TODO(bt,2025-04-24): Better error messages to caller + generate_image(request, &app_data_root, &sora_creds_manager) .await .map_err(|err| { @@ -82,7 +84,6 @@ pub async fn generate_image( // TODO(bt,2025-04-21): Read from in-memory cache instead, but allow for desktop replacement. let mut creds = read_sora_credentials_from_disk(app_data_root) - .await .map_err(|err| { error!("Failed to read Sora credentials from disk: {:?}", err); err diff --git a/crates/desktop/tauri-artcraft/src/lib.rs b/crates/desktop/tauri-artcraft/src/lib.rs index bd5bf94858..9624d673e4 100644 --- a/crates/desktop/tauri-artcraft/src/lib.rs +++ b/crates/desktop/tauri-artcraft/src/lib.rs @@ -28,7 +28,7 @@ pub async fn run() { let app_data_root2 = config.app_data_root.clone(); println!("Attempting to read existing credentials..."); - let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root).await; + let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root); let sora_creds_manager2 = sora_creds_manager.clone(); println!("Initializing backend runtime..."); diff --git a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs index 0726ea796b..2038c13e56 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs @@ -1,3 +1,4 @@ pub mod read_sora_credentials_from_disk; +pub mod read_sora_credentials_from_disk_with_upgrade; 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/read_sora_credentials_from_disk.rs b/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs index d5d3d22ad9..9d4579249c 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs @@ -1,16 +1,13 @@ use crate::state::app_dir::AppDataRoot; use anyhow::anyhow; use errors::AnyhowResult; -use log::{info, warn}; use openai_sora_client::creds::sora_cookies::SoraCookies; use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; use openai_sora_client::creds::sora_jwt_bearer_token::SoraJwtBearerToken; use openai_sora_client::creds::sora_sentinel::SoraSentinel; -use openai_sora_client::recipes::upgrade_or_renew_session::upgrade_or_renew_session; -use std::fs::{read_to_string, OpenOptions}; -use std::io::Write; +use std::fs::read_to_string; -pub async fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> AnyhowResult { +pub fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> AnyhowResult { let cookie_file = app_data_root.get_sora_cookie_file_path(); let bearer_file = app_data_root.get_sora_bearer_token_file_path(); let sentinel_file = app_data_root.get_sora_sentinel_file_path(); @@ -43,50 +40,9 @@ pub async fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> Any sentinel = Some(SoraSentinel::new(value)); } - let mut credentials = SoraCredentialSet::initialize( + Ok(SoraCredentialSet::initialize( cookie, bearer, sentinel, - ); - - let response = upgrade_or_renew_session(&mut credentials).await; - - if let Err(err) = response { - // NB: Make this portion infallible. Don't die on setup. We can recover downstream. - warn!("Failed to upgrade or renew session: {:?}", err); - } - - if !bearer_file.exists() { - if let Some(bearer) = &credentials.jwt_bearer_token { - info!("Persisting bearer token value..."); - let value = bearer.token_str(); - - let mut file = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&bearer_file)?; - - file.write_all(value.as_bytes())?; - file.flush()?; - } - } - - if !sentinel_file.exists() { - if let Some(sentinel) = &credentials.sora_sentinel { - info!("Persisting sentinel value..."); - let value = sentinel.get_sentinel(); - - let mut file = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&sentinel_file)?; - - file.write_all(value.as_bytes())?; - file.flush()?; - } - } - - Ok(credentials) + )) } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk_with_upgrade.rs b/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk_with_upgrade.rs new file mode 100644 index 0000000000..20e89f6d3b --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk_with_upgrade.rs @@ -0,0 +1,92 @@ +use crate::state::app_dir::AppDataRoot; +use anyhow::anyhow; +use errors::AnyhowResult; +use log::{info, warn}; +use openai_sora_client::creds::sora_cookies::SoraCookies; +use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; +use openai_sora_client::creds::sora_jwt_bearer_token::SoraJwtBearerToken; +use openai_sora_client::creds::sora_sentinel::SoraSentinel; +use openai_sora_client::recipes::upgrade_or_renew_session::upgrade_or_renew_session; +use std::fs::{read_to_string, OpenOptions}; +use std::io::Write; + +pub async fn read_sora_credentials_from_disk_with_upgrade(app_data_root: &AppDataRoot) -> AnyhowResult { + let cookie_file = app_data_root.get_sora_cookie_file_path(); + let bearer_file = app_data_root.get_sora_bearer_token_file_path(); + let sentinel_file = app_data_root.get_sora_sentinel_file_path(); + + if !cookie_file.exists() { + return Err(anyhow!("Cookie file does not exist: {:?}", cookie_file)); + } + + let value = read_to_string(cookie_file)? + .trim() + .to_string(); + + let cookie = SoraCookies::new(value); + + let mut bearer = None; + let mut sentinel = None; + + if bearer_file.exists() { + let value = read_to_string(&bearer_file)? + .trim() + .to_string(); + bearer = Some(SoraJwtBearerToken::new(value)?); + } + + if sentinel_file.exists() { + let value = read_to_string(&sentinel_file)? + .trim() + .to_string(); + + sentinel = Some(SoraSentinel::new(value)); + } + + let mut credentials = SoraCredentialSet::initialize( + cookie, + bearer, + sentinel, + ); + + let response = upgrade_or_renew_session(&mut credentials).await; + + if let Err(err) = response { + // NB: Make this portion infallible. Don't die on setup. We can recover downstream. + warn!("Failed to upgrade or renew session: {:?}", err); + } + + if !bearer_file.exists() { + if let Some(bearer) = &credentials.jwt_bearer_token { + info!("Persisting bearer token value..."); + let value = bearer.token_str(); + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&bearer_file)?; + + file.write_all(value.as_bytes())?; + file.flush()?; + } + } + + if !sentinel_file.exists() { + if let Some(sentinel) = &credentials.sora_sentinel { + info!("Persisting sentinel value..."); + let value = sentinel.get_sentinel(); + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&sentinel_file)?; + + file.write_all(value.as_bytes())?; + file.flush()?; + } + } + + Ok(credentials) +} 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 index 4808e1ef03..28ea7adece 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs @@ -19,10 +19,10 @@ pub struct SoraCredentialManager { impl SoraCredentialManager { - pub async fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { + pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { let holder = SoraCredentialHolder::new(); - match read_sora_credentials_from_disk(app_data_root).await { + 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"); @@ -64,8 +64,8 @@ impl SoraCredentialManager { self.holder.get_credentials_required() } - pub async fn reset_from_disk(&self) -> AnyhowResult { - let creds = read_sora_credentials_from_disk(&self.app_data_root).await?; + 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) } 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 3cad648c9e..c0b9597587 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 @@ -1,3 +1,4 @@ +use std::fs; use crate::state::app_dir::AppDataRoot; use crate::state::sora::sora_credential_holder::SoraCredentialHolder; use crate::state::sora::sora_credential_manager::SoraCredentialManager; @@ -10,7 +11,8 @@ use once_cell::sync::Lazy; use reqwest::Url; use std::fs::{read_to_string, OpenOptions}; use std::io::Write; -use chrono::{DateTime, NaiveDateTime}; +use std::ops::Sub; +use chrono::{DateTime, NaiveDateTime, TimeDelta}; use tauri::{AppHandle, Manager, Webview}; pub const LOGIN_WINDOW_NAME: &str = "login_window"; @@ -35,7 +37,11 @@ pub async fn sora_session_login_thread( loop { for (window_name, webview) in app.webviews() { if window_name == LOGIN_WINDOW_NAME { - let result = check_login_window(&webview, &app_data_root).await; + let result = check_login_window( + &webview, + &app_data_root, + &sora_creds_manager, + ).await; if let Err(err) = result { error!("Error checking login window: {:?}", err); } @@ -46,10 +52,14 @@ pub async fn sora_session_login_thread( } } -async fn check_login_window(webview: &Webview, app_data_root: &AppDataRoot) -> AnyhowResult<()> { +async fn check_login_window( + webview: &Webview, + app_data_root: &AppDataRoot, + sora_credential_manager: &SoraCredentialManager, +) -> AnyhowResult<()> { clear_browsing_data_on_test_domain(webview)?; //keep_on_task(webview)?; - extract_cookies_to_file(webview, app_data_root)?; + extract_cookies_to_file(webview, app_data_root, sora_credential_manager)?; //initialize_sora_jwt_bearer_token(app_data_root).await?; // TODO: This only runs once. We need better management. Ok(()) } @@ -93,7 +103,11 @@ fn clear_browsing_data_on_test_domain(webview: &Webview) -> AnyhowResult<()> { } // TODO(bt,2025-04-07): Heuristic to detect when logged in. Only write when logged in. -fn extract_cookies_to_file(webview: &Webview, app_data_root: &AppDataRoot) -> AnyhowResult<()> { +fn extract_cookies_to_file( + webview: &Webview, + app_data_root: &AppDataRoot, + sora_credential_manager: &SoraCredentialManager, +) -> AnyhowResult<()> { let new_cookies = get_all_sora_cookies_as_string(webview)?.trim().to_string(); let maybe_old_cookies = read_sora_cookies_from_disk(app_data_root); @@ -103,6 +117,18 @@ fn extract_cookies_to_file(webview: &Webview, app_data_root: &AppDataRoot) -> An should_write_cookies = old_cookies != new_cookies; } + if !should_write_cookies { + return Ok(()); + } + + // TODO(bt): Race conditions ahead. + + sora_credential_manager.clear_credentials()?; + + fs::remove_file(app_data_root.get_sora_cookie_file_path())?; + fs::remove_file(app_data_root.get_sora_bearer_token_file_path())?; + fs::remove_file(app_data_root.get_sora_sentinel_file_path())?; + let cookie_file = app_data_root.get_sora_cookie_file_path(); let mut file = OpenOptions::new() @@ -114,6 +140,8 @@ fn extract_cookies_to_file(webview: &Webview, app_data_root: &AppDataRoot) -> An file.write_all(new_cookies.as_bytes())?; file.flush()?; + sora_credential_manager.reset_from_disk()?; + Ok(()) }