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 785456410d..5db4de9e99 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 @@ -9,6 +9,8 @@ use openai_sora_client::creds::sora_sentinel::SoraSentinel; use openai_sora_client::requests::sentinel_refresh::generate::token::generate_token; use std::fs::OpenOptions; use std::io::Write; +use openai_sora_client::creds::sora_cookies::SoraCookies; +use openai_sora_client::creds::sora_jwt_bearer_token::SoraJwtBearerToken; #[derive(Clone)] pub struct SoraCredentialManager { @@ -57,6 +59,25 @@ impl SoraCredentialManager { Ok(creds) } + pub fn persist_all_to_disk(&self) -> AnyhowResult<()> { + let credentials = self.holder.get_credentials_required()?; + + info!("Persisting cookies to disk..."); + persist_cookies_to_disk(&credentials.cookies, &self.app_data_root)?; + + if let Some(bearer) = &credentials.jwt_bearer_token { + info!("Persisting JWT to disk..."); + persist_jwt_bearer_to_disk(bearer, &self.app_data_root)?; + } + + if let Some(sentinel) = &credentials.sora_sentinel { + info!("Persisting sentinel to disk..."); + persist_sentinel_to_disk(sentinel, &self.app_data_root)?; + } + + Ok(()) + } + /// Refresh the sentinel token from Sora's API pub async fn call_sentinel_refresh(&self) -> AnyhowResult { // NB(bt,2025-04-21): Technically we don't need credentials to get a sentinel. @@ -89,3 +110,48 @@ impl SoraCredentialManager { Ok(creds) } } + +fn persist_cookies_to_disk(cookies: &SoraCookies, app_data_root: &AppDataRoot) -> AnyhowResult<()> { + let filename = app_data_root.get_sora_cookie_file_path(); + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(filename)?; + + file.write_all(cookies.as_bytes())?; + file.flush()?; + + Ok(()) +} + +fn persist_jwt_bearer_to_disk(bearer: &SoraJwtBearerToken, app_data_root: &AppDataRoot) -> AnyhowResult<()> { + let filename = app_data_root.get_sora_bearer_token_file_path(); + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(filename)?; + + file.write_all(bearer.as_bytes())?; + file.flush()?; + + Ok(()) +} + +fn persist_sentinel_to_disk(sentinel: &SoraSentinel, app_data_root: &AppDataRoot) -> AnyhowResult<()> { + let filename = app_data_root.get_sora_sentinel_file_path(); + + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(filename)?; + + file.write_all(sentinel.as_bytes())?; + file.flush()?; + + 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 391846714c..a68e102855 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 @@ -14,6 +14,9 @@ use std::io::Write; use std::ops::Sub; use chrono::{DateTime, NaiveDateTime, TimeDelta}; use tauri::{AppHandle, Manager, Webview}; +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}; pub const LOGIN_WINDOW_NAME: &str = "login_window"; @@ -59,7 +62,7 @@ async fn check_login_window( ) -> AnyhowResult<()> { clear_browsing_data_on_test_domain(webview)?; //keep_on_task(webview)?; - extract_cookies_to_file(webview, app_data_root, sora_credential_manager)?; + extract_cookies_to_file(webview, app_data_root, sora_credential_manager).await?; //initialize_sora_jwt_bearer_token(app_data_root).await?; // TODO: This only runs once. We need better management. Ok(()) } @@ -103,7 +106,7 @@ 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( +async fn extract_cookies_to_file( webview: &Webview, app_data_root: &AppDataRoot, sora_credential_manager: &SoraCredentialManager, @@ -114,15 +117,19 @@ fn extract_cookies_to_file( return Ok(()); } - let maybe_old_cookies = read_sora_cookies_from_disk(app_data_root); + let session_cookie_presence = has_session_cookie(&new_cookies) + .unwrap_or_else(|err| { + error!("Failed to check for session cookie: {:?}", err); + SessionCookiePresence::MaybePresent + }); - let mut should_write_cookies = maybe_old_cookies.is_none(); + info!("Session cookies are: {:?}", session_cookie_presence); - if new_cookies.len() > 50 { - if let Some(old_cookies) = maybe_old_cookies { - should_write_cookies = old_cookies != new_cookies; - } - } + let (should_write_cookies, should_upgrade_session) = match session_cookie_presence { + SessionCookiePresence::Present => (true, true), + SessionCookiePresence::MaybePresent => (true, false), + SessionCookiePresence::Absent => (false, false), + }; if !should_write_cookies { return Ok(()); @@ -136,18 +143,15 @@ fn extract_cookies_to_file( let _r = fs::remove_file(app_data_root.get_sora_bearer_token_file_path()); let _r = fs::remove_file(app_data_root.get_sora_sentinel_file_path()); - let cookie_file = app_data_root.get_sora_cookie_file_path(); + let mut new_credentials = + SoraCredentialSet::initialize_with_just_cookies_str(&new_cookies); - let mut file = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(cookie_file)?; + if should_upgrade_session { + let _upgraded = maybe_upgrade_or_renew_session(&mut new_credentials).await?; + } - file.write_all(new_cookies.as_bytes())?; - file.flush()?; - - sora_credential_manager.reset_from_disk()?; + sora_credential_manager.set_credentials(&new_credentials)?; + sora_credential_manager.persist_all_to_disk()?; Ok(()) } diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs index da4b02fd4a..5ff387b645 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs @@ -15,6 +15,10 @@ impl SoraCookies { &self.cookies } + pub fn as_bytes(&self) -> &[u8] { + self.cookies.as_bytes() + } + pub fn to_string(&self) -> String { self.cookies.clone() } diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs index 8176302f82..518f702d7f 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs @@ -21,6 +21,10 @@ impl SoraJwtBearerToken { &self.token } + pub fn as_bytes(&self) -> &[u8] { + self.token.as_bytes() + } + pub fn jwt_claims(&self) -> &SoraJwtClaims { &self.jwt_claims } diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs index 40682751ff..da9e492801 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs @@ -16,4 +16,8 @@ impl SoraSentinel { pub fn get_sentinel(&self) -> &str { &self.sentinel } + + pub fn as_bytes(&self) -> &[u8] { + self.sentinel.as_bytes() + } }