From 2eefaaf7a15cf0c215d0597761a3a3e427b52abc Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 30 Apr 2025 22:30:59 -0400 Subject: [PATCH] storyteller cookie saving; temporarily broke window size persist --- Cargo.lock | 1 + crates/desktop/tauri-artcraft/src/lib.rs | 11 +++- .../read_storyteller_credentials_from_disk.rs | 2 +- .../main_window_thread.rs | 30 +++++----- .../src/threads/main_window_thread/mod.rs | 3 + .../persist_storyteller_cookies_task.rs | 59 +++++++++++++++++++ .../persist_window_resize_task.rs | 26 ++++++++ .../src/threads/sora_session_login_thread.rs | 2 +- .../tauri-artcraft/src/utils/cookies/mod.rs | 2 + .../{ => cookies}/sora_webview_cookies.rs | 0 .../cookies/storyteller_webview_cookies.rs | 36 +++++++++++ .../desktop/tauri-artcraft/src/utils/mod.rs | 4 +- .../api_clients/storyteller_client/Cargo.toml | 1 + .../src/credentials/storyteller_avt_cookie.rs | 15 +++++ .../credentials/storyteller_credential_set.rs | 40 +++++++++++-- .../credentials/storyteller_session_cookie.rs | 15 +++++ 16 files changed, 220 insertions(+), 27 deletions(-) rename crates/desktop/tauri-artcraft/src/threads/{ => main_window_thread}/main_window_thread.rs (50%) create mode 100644 crates/desktop/tauri-artcraft/src/threads/main_window_thread/mod.rs create mode 100644 crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_storyteller_cookies_task.rs create mode 100644 crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_window_resize_task.rs create mode 100644 crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs rename crates/desktop/tauri-artcraft/src/utils/{ => cookies}/sora_webview_cookies.rs (100%) create mode 100644 crates/desktop/tauri-artcraft/src/utils/cookies/storyteller_webview_cookies.rs diff --git a/Cargo.lock b/Cargo.lock index d9ff358e50..b5da03bc2c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11280,6 +11280,7 @@ dependencies = [ "anyhow", "backoff", "chrono", + "cookie 0.18.1", "enums", "enums_public", "errors", diff --git a/crates/desktop/tauri-artcraft/src/lib.rs b/crates/desktop/tauri-artcraft/src/lib.rs index da07da9840..8543c3fb72 100644 --- a/crates/desktop/tauri-artcraft/src/lib.rs +++ b/crates/desktop/tauri-artcraft/src/lib.rs @@ -4,7 +4,7 @@ pub mod state; pub mod threads; pub mod utils; -use tauri::{Manager}; +use tauri::Manager; use crate::commands::flip_image::flip_image; use crate::commands::platform_info_command::platform_info_command; @@ -15,8 +15,9 @@ use crate::state::app_config::AppConfig; use crate::state::main_window_size::MainWindowSize; use crate::state::sora::sora_credential_manager::SoraCredentialManager; use crate::state::sora::sora_task_queue::SoraTaskQueue; +use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; use crate::threads::discord_presence_thread::discord_presence_thread; -use crate::threads::main_window_thread::main_window_thread; +use crate::threads::main_window_thread::main_window_thread::main_window_thread; use crate::threads::sora_session_login_thread::sora_session_login_thread; use crate::threads::sora_task_polling_thread::sora_task_polling_thread; @@ -35,6 +36,10 @@ pub fn run() { let app_data_root = config.app_data_root.clone(); let app_data_root_2 = config.app_data_root.clone(); + println!("Attempting to read existing artcraft credentials..."); + 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..."); let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root); let sora_creds_manager_2 = sora_creds_manager.clone(); @@ -89,7 +94,7 @@ pub fn run() { } tauri::async_runtime::spawn(sora_session_login_thread(app_2, app_data_root_2, sora_creds_manager_2)); - tauri::async_runtime::spawn(main_window_thread(app_3, app_data_root_3)); + tauri::async_runtime::spawn(main_window_thread(app_3, app_data_root_3, storyteller_creds_manager_2)); tauri::async_runtime::spawn(sora_task_polling_thread(app_data_root_4, sora_creds_manager_3, sora_task_queue_2)); tauri::async_runtime::spawn(discord_presence_thread()); diff --git a/crates/desktop/tauri-artcraft/src/state/storyteller/read_storyteller_credentials_from_disk.rs b/crates/desktop/tauri-artcraft/src/state/storyteller/read_storyteller_credentials_from_disk.rs index 842bfd717b..a721ec05b9 100644 --- a/crates/desktop/tauri-artcraft/src/state/storyteller/read_storyteller_credentials_from_disk.rs +++ b/crates/desktop/tauri-artcraft/src/state/storyteller/read_storyteller_credentials_from_disk.rs @@ -33,7 +33,7 @@ pub fn read_storyteller_credentials_from_disk(app_data_root: &AppDataRoot) -> An } Ok(StorytellerCredentialSet::initialize( - session, avt, + session, )) } diff --git a/crates/desktop/tauri-artcraft/src/threads/main_window_thread.rs b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/main_window_thread.rs similarity index 50% rename from crates/desktop/tauri-artcraft/src/threads/main_window_thread.rs rename to crates/desktop/tauri-artcraft/src/threads/main_window_thread/main_window_thread.rs index da93aa1bdc..f68941ae28 100644 --- a/crates/desktop/tauri-artcraft/src/threads/main_window_thread.rs +++ b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/main_window_thread.rs @@ -1,5 +1,8 @@ use crate::state::data_dir::app_data_root::AppDataRoot; use crate::state::main_window_size::MainWindowSize; +use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; +use crate::threads::main_window_thread::persist_storyteller_cookies_task::persist_storyteller_cookies_task; +use crate::threads::main_window_thread::persist_window_resize_task::persist_window_resize_task; use errors::AnyhowResult; use log::{error, info}; use tauri::{AppHandle, Manager, Webview, Window}; @@ -9,6 +12,7 @@ const MAIN_WINDOW_NAME : &str = "main"; pub async fn main_window_thread( app: AppHandle, app_data_root: AppDataRoot, + storyteller_creds_manager: StorytellerCredentialManager, ) -> ! { loop { for (window_name, window) in app.windows() { @@ -16,6 +20,7 @@ pub async fn main_window_thread( let result = handle_main_window( &window, &app_data_root, + &storyteller_creds_manager, ).await; if let Err(err) = result { error!("Error handling main window: {:?}", err); @@ -29,20 +34,17 @@ pub async fn main_window_thread( pub async fn handle_main_window( window: &Window, app_data_root: &AppDataRoot, + storyteller_creds_manager: &StorytellerCredentialManager, ) -> AnyhowResult<()> { - let mut window_size = MainWindowSize::from_window(window)?; - loop { - let old_size = window_size.to_physical_size(); - let new_size = window.inner_size()?; - - if !window_size.matches_physical_size(&new_size) { - info!("Window size changed from {:?} to {:?}", old_size, new_size); - window_size = MainWindowSize::from_window(&window)?; - info!("Saving window size configs to disk..."); - window_size.persist_to_filesystem(app_data_root)?; - } - - tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + log_errors(persist_storyteller_cookies_task(window, app_data_root, storyteller_creds_manager).await); + log_errors(persist_window_resize_task(window, app_data_root).await); + tokio::time::sleep(std::time::Duration::from_millis(1_000)).await; } -} \ No newline at end of file +} + +pub fn log_errors(result: AnyhowResult) { + if let Err(err) = result { + error!("Error persisting window size: {:?}", err); + } +} diff --git a/crates/desktop/tauri-artcraft/src/threads/main_window_thread/mod.rs b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/mod.rs new file mode 100644 index 0000000000..635ce85c2e --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/mod.rs @@ -0,0 +1,3 @@ +pub mod main_window_thread; +pub mod persist_storyteller_cookies_task; +pub mod persist_window_resize_task; \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_storyteller_cookies_task.rs b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_storyteller_cookies_task.rs new file mode 100644 index 0000000000..91951ba1f5 --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_storyteller_cookies_task.rs @@ -0,0 +1,59 @@ +use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; +use crate::threads::sora_session_login_thread::LOGIN_WINDOW_NAME; +use crate::utils::cookies::storyteller_webview_cookies::get_storyteller_cookies; +use anyhow::anyhow; +use errors::AnyhowResult; +use log::{error, info}; +use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; +use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session; +use openai_sora_client::utils::has_session_cookie::{has_session_cookie, SessionCookiePresence}; +use std::fs; +use tauri::{AppHandle, Manager, Webview, Window}; + +const MAIN_WEBVIEW_NAME: &str = "main"; + +pub async fn persist_storyteller_cookies_task( + window: &Window, + app_data_root: &AppDataRoot, + storyteller_credential_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + for webview in window.webviews() { + let label = webview.label(); + info!("Webview label: {:?}", label); + if label == MAIN_WEBVIEW_NAME { + persist_webview_cookies(&webview, app_data_root, storyteller_credential_manager).await?; + break; + } + } + Ok(()) +} +async fn persist_webview_cookies( + webview: &Webview, + app_data_root: &AppDataRoot, + storyteller_credential_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + let webview_credentials = get_storyteller_cookies(webview)?; + + if webview_credentials.is_empty() { + // TODO: handle logout + return Ok(()); + } + + let mut replace_credentials = true; + + let maybe_old_credentials = storyteller_credential_manager.get_credentials()?; + + if let Some(old_credentials) = maybe_old_credentials { + if old_credentials.equals(&webview_credentials) { + replace_credentials = false; + } + } + + if replace_credentials { + storyteller_credential_manager.set_credentials(&webview_credentials)?; + storyteller_credential_manager.persist_all_to_disk()?; + } + + Ok(()) +} diff --git a/crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_window_resize_task.rs b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_window_resize_task.rs new file mode 100644 index 0000000000..fcea642c70 --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/persist_window_resize_task.rs @@ -0,0 +1,26 @@ +use log::info; +use tauri::Window; +use errors::AnyhowResult; +use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::main_window_size::MainWindowSize; + +pub async fn persist_window_resize_task( + window: &Window, + app_data_root: &AppDataRoot, +) -> AnyhowResult<()> { + let mut window_size = MainWindowSize::from_window(window)?; + + let new_size = window.inner_size()?; + + // TODO: Temporary regression. This used to work, but we no longer have old state to compare against. + // if !window_size.matches_physical_size(&new_size) { + // let old_size = window_size.to_physical_size(); + // info!("Window size changed from {:?} to {:?}", old_size, new_size); + // window_size = MainWindowSize::from_window(&window)?; + // info!("Saving window size configs to disk..."); + // window_size.persist_to_filesystem(app_data_root)?; + // } + + Ok(()) +} + 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 e0c9a02bf0..69c4a40aa5 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,8 +1,8 @@ use crate::state::data_dir::app_data_root::AppDataRoot; use crate::state::sora::sora_credential_holder::SoraCredentialHolder; use crate::state::sora::sora_credential_manager::SoraCredentialManager; +use crate::utils::cookies::sora_webview_cookies::get_all_sora_cookies_as_string; use crate::utils::sora::initialize_sora_jwt_bearer_token::initialize_sora_jwt_bearer_token; -use crate::utils::sora_webview_cookies::get_all_sora_cookies_as_string; use anyhow::anyhow; use chrono::{DateTime, NaiveDateTime, TimeDelta}; use errors::AnyhowResult; diff --git a/crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs b/crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs new file mode 100644 index 0000000000..c1b8ecbb42 --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs @@ -0,0 +1,2 @@ +pub mod sora_webview_cookies; +pub mod storyteller_webview_cookies; \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/utils/sora_webview_cookies.rs b/crates/desktop/tauri-artcraft/src/utils/cookies/sora_webview_cookies.rs similarity index 100% rename from crates/desktop/tauri-artcraft/src/utils/sora_webview_cookies.rs rename to crates/desktop/tauri-artcraft/src/utils/cookies/sora_webview_cookies.rs diff --git a/crates/desktop/tauri-artcraft/src/utils/cookies/storyteller_webview_cookies.rs b/crates/desktop/tauri-artcraft/src/utils/cookies/storyteller_webview_cookies.rs new file mode 100644 index 0000000000..3258803d82 --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/utils/cookies/storyteller_webview_cookies.rs @@ -0,0 +1,36 @@ +use errors::AnyhowResult; +use once_cell::sync::Lazy; +use reqwest::Url; +use storyteller_client::credentials::storyteller_avt_cookie::StorytellerAvtCookie; +use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use storyteller_client::credentials::storyteller_session_cookie::StorytellerSessionCookie; +use tauri::Webview; + +const STORYTELLER_ROOT_COOKIE_URL_STR: &str = "https://api.storyteller.ai/"; + +static STORYTELLER_ROOT_COOKIE_URL: Lazy = Lazy::new(|| { + Url::parse(STORYTELLER_ROOT_COOKIE_URL_STR).expect("URL should parse") +}); + +const AVT_COOKIE_NAME : &str = "visitor"; +const SESSION_COOKIE_NAME : &str = "session"; + +pub fn get_storyteller_cookies(webview: &Webview) -> AnyhowResult { + let cookies = webview.cookies_for_url(STORYTELLER_ROOT_COOKIE_URL.clone())?; + + let mut avt_cookie = None; + let mut session_cookie = None; + + for cookie in cookies.into_iter() { + if let Some(avt) = StorytellerAvtCookie::maybe_from_cookie(&cookie) { + avt_cookie = Some(avt); + } else if let Some(session) = StorytellerSessionCookie::maybe_from_cookie(&cookie) { + session_cookie = Some(session); + } + } + + Ok(StorytellerCredentialSet::initialize( + avt_cookie, + session_cookie, + )) +} diff --git a/crates/desktop/tauri-artcraft/src/utils/mod.rs b/crates/desktop/tauri-artcraft/src/utils/mod.rs index a685e469e6..65f49bffd7 100644 --- a/crates/desktop/tauri-artcraft/src/utils/mod.rs +++ b/crates/desktop/tauri-artcraft/src/utils/mod.rs @@ -1,7 +1,7 @@ pub mod clear_all_webview_cookies; +pub mod cookies; pub mod get_url_file_extension; pub mod image; pub mod large_file_transfer; pub mod simple_http_download; -pub mod sora; -pub mod sora_webview_cookies; \ No newline at end of file +pub mod sora; \ No newline at end of file diff --git a/crates/lib/api_clients/storyteller_client/Cargo.toml b/crates/lib/api_clients/storyteller_client/Cargo.toml index a68f7bf416..b4155c50e3 100644 --- a/crates/lib/api_clients/storyteller_client/Cargo.toml +++ b/crates/lib/api_clients/storyteller_client/Cargo.toml @@ -25,6 +25,7 @@ tokens = { workspace = true } # External anyhow = { workspace = true } chrono = { workspace = true } +cookie = { workspace = true } log = { workspace = true } reqwest = { workspace = true, features = ["multipart", "gzip"] } url = { workspace = true } diff --git a/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_avt_cookie.rs b/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_avt_cookie.rs index a4b13dc276..a836288491 100644 --- a/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_avt_cookie.rs +++ b/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_avt_cookie.rs @@ -1,3 +1,6 @@ +use cookie::Cookie; + +const AVT_COOKIE_NAME : &str = "visitor"; /// The "anonymous visitor token" cookie allows users to interact with /// the service in a logged-out state with some continuity. @@ -11,6 +14,14 @@ impl StorytellerAvtCookie { Self { cookie } } + pub fn maybe_from_cookie(cookie: &Cookie) -> Option { + if cookie.name() == AVT_COOKIE_NAME { + Some(Self::new(cookie.value().to_string())) + } else { + None + } + } + pub fn as_str(&self) -> &str { &self.cookie } @@ -22,4 +33,8 @@ impl StorytellerAvtCookie { pub fn to_string(&self) -> String { self.cookie.clone() } + + pub fn equals(&self, other: &Self) -> bool { + self.cookie == other.cookie + } } diff --git a/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_credential_set.rs b/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_credential_set.rs index d22d7cacd6..6c6a407648 100644 --- a/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_credential_set.rs +++ b/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_credential_set.rs @@ -3,39 +3,67 @@ use crate::credentials::storyteller_session_cookie::StorytellerSessionCookie; #[derive(Clone)] pub struct StorytellerCredentialSet { - pub session: Option, pub avt: Option, + pub session: Option, } impl StorytellerCredentialSet { pub fn empty() -> Self { Self { - session: None, avt: None, + session: None, } } pub fn initialize( - session: Option, avt: Option, + session: Option, ) -> Self { Self { - session, avt, + session, } } pub fn initialize_with_just_cookie(session: StorytellerSessionCookie) -> Self { Self { - session: Some(session), avt: None, + session: Some(session), } } pub fn initialize_with_just_avt(avt: StorytellerAvtCookie) -> Self { Self { - session: None, avt: Some(avt), + session: None, } } + + pub fn is_empty(&self) -> bool { + self.avt.is_none() && self.session.is_none() + } + + pub fn equals(&self, other: &Self) -> bool { + match (&self.avt, &other.avt) { + (None, None) => {} // Fallthrough + (Some(avt), Some(other_avt)) => { + if !avt.equals(other_avt) { + return false; + } + } + _ => return false, + } + + match (&self.session, &other.session) { + (None, None) => {} // Fallthrough + (Some(session), Some(other_session)) => { + if !session.equals(other_session) { + return false; + } + } + _ => return false, + } + + true + } } diff --git a/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_session_cookie.rs b/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_session_cookie.rs index 638605f941..58cd6e6dd7 100644 --- a/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_session_cookie.rs +++ b/crates/lib/api_clients/storyteller_client/src/credentials/storyteller_session_cookie.rs @@ -1,3 +1,6 @@ +use cookie::Cookie; + +const SESSION_COOKIE_NAME : &str = "session"; /// The primary FakeYou / Storyteller / Artcraft backend session identifier. #[derive(Clone)] @@ -9,6 +12,14 @@ impl StorytellerSessionCookie { pub fn new(cookie: String) -> Self { Self { cookie } } + + pub fn maybe_from_cookie(cookie: &Cookie) -> Option { + if cookie.name() == SESSION_COOKIE_NAME { + Some(Self::new(cookie.value().to_string())) + } else { + None + } + } pub fn as_str(&self) -> &str { &self.cookie @@ -21,4 +32,8 @@ impl StorytellerSessionCookie { pub fn to_string(&self) -> String { self.cookie.clone() } + + pub fn equals(&self, other: &Self) -> bool { + self.cookie == other.cookie + } }