From 8a364605cf85f03aaff11e806ee4b4fc073eda64 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 1 May 2025 01:10:45 -0400 Subject: [PATCH] bug is still present, might be a race condition --- crates/desktop/tauri-artcraft/src/lib.rs | 2 +- .../main_window_thread/main_window_thread.rs | 4 +- .../persist_storyteller_cookies_task.rs | 53 ++++++++++++++++--- .../tauri-artcraft/src/utils/cookies/mod.rs | 3 +- .../cookies/storyteller_webview_cookies.rs | 36 ------------- 5 files changed, 49 insertions(+), 49 deletions(-) delete mode 100644 crates/desktop/tauri-artcraft/src/utils/cookies/storyteller_webview_cookies.rs diff --git a/crates/desktop/tauri-artcraft/src/lib.rs b/crates/desktop/tauri-artcraft/src/lib.rs index 8ff17e0d7f..a5489d9926 100644 --- a/crates/desktop/tauri-artcraft/src/lib.rs +++ b/crates/desktop/tauri-artcraft/src/lib.rs @@ -94,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, storyteller_creds_manager_2)); + 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/threads/main_window_thread/main_window_thread.rs b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/main_window_thread.rs index ddffb7fe95..7e8495b479 100644 --- a/crates/desktop/tauri-artcraft/src/threads/main_window_thread/main_window_thread.rs +++ b/crates/desktop/tauri-artcraft/src/threads/main_window_thread/main_window_thread.rs @@ -16,7 +16,7 @@ pub async fn main_window_thread( storyteller_creds_manager: StorytellerCredentialManager, ) -> ! { let window_size_slot: CloneSlot = CloneSlot::empty(); - + loop { for (window_name, window) in app.windows() { if window_name == MAIN_WINDOW_NAME { @@ -42,8 +42,8 @@ pub async fn handle_main_window( window_size_slot: &CloneSlot, ) -> AnyhowResult<()> { loop { - log_errors(persist_storyteller_cookies_task(window, app_data_root, storyteller_creds_manager).await); log_errors(persist_window_resize_task(window, app_data_root, window_size_slot).await); + log_errors(persist_storyteller_cookies_task(window, app_data_root, storyteller_creds_manager).await); tokio::time::sleep(std::time::Duration::from_millis(1_000)).await; } } 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 index 051c534cd7..624ec5c8c9 100644 --- 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 @@ -1,18 +1,28 @@ 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 log::{debug, error, info}; +use once_cell::sync::Lazy; +use reqwest::Url; use std::fs; +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::{AppHandle, Manager, Webview, Window}; const MAIN_WEBVIEW_NAME: &str = "main"; +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 async fn persist_storyteller_cookies_task( window: &Window, app_data_root: &AppDataRoot, @@ -20,7 +30,6 @@ pub async fn persist_storyteller_cookies_task( ) -> 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; @@ -33,7 +42,7 @@ async fn persist_webview_cookies( app_data_root: &AppDataRoot, storyteller_credential_manager: &StorytellerCredentialManager, ) -> AnyhowResult<()> { - let current_webview_credentials = get_storyteller_cookies(webview)?; + let current_webview_credentials = get_storyteller_cookies(webview).await?; if current_webview_credentials.is_empty() { // TODO: handle logout / cookie deletion @@ -41,7 +50,7 @@ async fn persist_webview_cookies( } let mut replace_credentials = true; - + let maybe_old_credentials = storyteller_credential_manager.get_credentials()?; if let Some(old_credentials) = maybe_old_credentials { @@ -58,3 +67,31 @@ async fn persist_webview_cookies( Ok(()) } + +pub async fn get_storyteller_cookies(webview: &Webview) -> AnyhowResult { + debug!("Getting storyteller cookies..."); + + // FIXME: THIS IS A TOTAL HACK. / DO NOT REMOVE. + // NB: If we don't call a method on the webview, the very next call to get cookies will deadlock. + if let Ok(url) = webview.url() { + debug!("For url: {:?}", url); + } + + 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/cookies/mod.rs b/crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs index c1b8ecbb42..fc9a63ecd8 100644 --- a/crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs +++ b/crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs @@ -1,2 +1 @@ -pub mod sora_webview_cookies; -pub mod storyteller_webview_cookies; \ No newline at end of file +pub mod sora_webview_cookies; \ No newline at end of file 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 deleted file mode 100644 index 3258803d82..0000000000 --- a/crates/desktop/tauri-artcraft/src/utils/cookies/storyteller_webview_cookies.rs +++ /dev/null @@ -1,36 +0,0 @@ -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, - )) -}