storyteller cookie saving; temporarily broke window size persist

Cargo.lock

@@ -11280,6 +11280,7 @@ dependencies = [
  "anyhow",
  "backoff",
  "chrono",
+ "cookie 0.18.1",
  "enums",
  "enums_public",
  "errors",

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());
 

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,
   ))
 }

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;
   }
-}
+}
+
+pub fn log_errors<T>(result: AnyhowResult<T>) {
+  if let Err(err) = result {
+    error!("Error persisting window size: {:?}", err);
+  }
+}

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;

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(())
+}

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(())
+}
+

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;

crates/desktop/tauri-artcraft/src/utils/cookies/mod.rs

@@ -0,0 +1,2 @@
+pub mod sora_webview_cookies;
+pub mod storyteller_webview_cookies;

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<Url> = 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<StorytellerCredentialSet> {
+  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,
+  ))
+}

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;
+pub mod sora;

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 }

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<Self> {
+    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
+  }
 }

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<StorytellerSessionCookie>,
   pub avt: Option<StorytellerAvtCookie>,
+  pub session: Option<StorytellerSessionCookie>,
 }
 
 impl StorytellerCredentialSet {
   pub fn empty() -> Self {
     Self {
-      session: None,
       avt: None,
+      session: None,
     }
   }
   
   pub fn initialize(
-    session: Option<StorytellerSessionCookie>,
     avt: Option<StorytellerAvtCookie>,
+    session: Option<StorytellerSessionCookie>,
   ) -> 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
+  }
 }

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<Self> {
+    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
+  }
 }