From d2ab68d1a1f0f7ea7367ce3f37d7e10092ca2a09 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 24 Apr 2025 08:55:21 -0400 Subject: [PATCH] load or request bearer/sentinel at start, more refactor --- .../sora/sora_image_generation_command.rs | 11 +- .../commands/sora/sora_image_remix_command.rs | 17 ++- crates/desktop/tauri-artcraft/src/lib.rs | 4 +- .../sora/read_sora_credentials_from_disk.rs | 109 ++++++++++++------ .../src/state/sora/sora_credential_manager.rs | 16 ++- .../src/creds/sora_credential_set.rs | 12 ++ 6 files changed, 109 insertions(+), 60 deletions(-) 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 131c5800a4..f90ee163a4 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 @@ -8,6 +8,7 @@ use image::codecs::png::{CompressionType, FilterType, PngEncoder}; use image::{DynamicImage, ImageReader}; use log::{error, info}; use openai_sora_client::credentials::SoraCredentials; +use openai_sora_client::creds::credential_migration::CredentialMigrationRef; use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages}; use openai_sora_client::requests::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest}; use openai_sora_client::requests::upload::upload_media_from_bytes::sora_media_upload_from_bytes; @@ -15,7 +16,6 @@ use openai_sora_client::requests::upload::upload_media_from_file::SoraMediaUploa use std::fs::read_to_string; use std::io::Cursor; use tauri::{AppHandle, Manager, State}; -use openai_sora_client::creds::credential_migration::CredentialMigrationRef; #[tauri::command] pub async fn sora_image_generation_command( @@ -44,14 +44,13 @@ pub async fn generate_image( sora_creds_holder: &SoraCredentialHolder, ) -> AnyhowResult<()> { - let sora_credentials_payload = read_sora_credentials_from_disk(app_data_root) + let creds = read_sora_credentials_from_disk(app_data_root) + .await .map_err(|err| { error!("Failed to read Sora credentials from disk: {:?}", err); err })?; - //let sora_credentials = sora_creds_holder.get_credentials()?; - let mut sora_media_tokens = vec![]; if let Some(image) = maybe_image { @@ -59,7 +58,7 @@ pub async fn generate_image( let filename = "image.png".to_string(); - let response = sora_media_upload_from_bytes(image_bytes, filename, CredentialMigrationRef::New(&sora_credentials_payload.credentials_set)) + let response = sora_media_upload_from_bytes(image_bytes, filename, CredentialMigrationRef::New(&creds)) .await .map_err(|err| { error!("Failed to upload image to Sora: {:?}", err); @@ -74,7 +73,7 @@ pub async fn generate_image( num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: CredentialMigrationRef::New(&sora_credentials_payload.credentials_set), + credentials: CredentialMigrationRef::New(&creds), }).await .map_err(|err| { error!("Failed to call Sora image generation: {:?}", 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 bbf6aa29bc..cece13ad45 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 @@ -81,21 +81,20 @@ pub async fn generate_image( let files_to_upload = vec![filename]; // TODO(bt,2025-04-21): Read from in-memory cache instead, but allow for desktop replacement. - let sora_creds_payload= read_sora_credentials_from_disk(app_data_root) + 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 })?; - let mut sora_creds = sora_creds_payload.credentials_set; - let no_sentinel = !app_data_root.get_sora_sentinel_file_path().is_file() - || sora_creds.sora_sentinel.is_none(); + || creds.sora_sentinel.is_none(); if no_sentinel { info!("Refreshing credentials..."); - sora_creds = sora_creds_manager.call_sentinel_refresh() + creds = sora_creds_manager.call_sentinel_refresh() .await .map_err(|err| { error!("Failed to refresh: {:?}", err); @@ -108,7 +107,7 @@ pub async fn generate_image( let mut sora_media_tokens = Vec::with_capacity(files_to_upload.len()); for file_path in files_to_upload { - let sora_upload_response = sora_media_upload_from_file(file_path, CredentialMigrationRef::New(&sora_creds)) + let sora_upload_response = sora_media_upload_from_file(file_path, CredentialMigrationRef::New(&creds)) .await?; sora_media_tokens.push(sora_upload_response.id); @@ -123,13 +122,13 @@ pub async fn generate_image( num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: CredentialMigrationRef::New(&sora_creds), + credentials: CredentialMigrationRef::New(&creds), }).await; if let Err(err) = &response { error!("Error in generating image: {:?}", err); - sora_creds = sora_creds_manager.call_sentinel_refresh() + creds = sora_creds_manager.call_sentinel_refresh() .await .map_err(|err| { error!("Failed to refresh: {:?}", err); @@ -143,7 +142,7 @@ pub async fn generate_image( num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: CredentialMigrationRef::New(&sora_creds), + credentials: CredentialMigrationRef::New(&creds), }).await; } diff --git a/crates/desktop/tauri-artcraft/src/lib.rs b/crates/desktop/tauri-artcraft/src/lib.rs index b08eaad4ae..bd5bf94858 100644 --- a/crates/desktop/tauri-artcraft/src/lib.rs +++ b/crates/desktop/tauri-artcraft/src/lib.rs @@ -16,7 +16,7 @@ use tauri_plugin_log::Target; use tauri_plugin_log::TargetKind; #[cfg_attr(mobile, tauri::mobile_entry_point)] -pub fn run() { +pub async fn run() { // NB: Tauri wants to install the logger itself, so we can't rely on the logger crate // until the tauri runtime begins. println!("Loading model config..."); @@ -28,7 +28,7 @@ pub 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); + let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root).await; let sora_creds_manager2 = sora_creds_manager.clone(); println!("Initializing backend runtime..."); 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 91591b819a..d5d3d22ad9 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,51 +1,92 @@ -use std::fs::read_to_string; -use errors::AnyhowResult; -use openai_sora_client::credentials::SoraCredentials; -use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; 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; -// TODO: This is only to migrate to the new type. -pub struct CredentialsPayload { - pub credentials: SoraCredentials, - pub credentials_set: SoraCredentialSet, -} - -pub fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> AnyhowResult { +pub async 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(); - let cookie = read_to_string(cookie_file)? + 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 bearer = read_to_string(bearer_file)? - .trim() - .to_string(); + let cookie = SoraCookies::new(value); - let mut credentials; + let mut bearer = None; + let mut sentinel = None; - if sentinel_file.exists() && sentinel_file.is_file() { - let sentinel = read_to_string(sentinel_file)? + 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(); - credentials = SoraCredentials { - bearer_token: bearer, - cookie, - sentinel: Some(sentinel), - }; - } else { - - credentials = SoraCredentials { - bearer_token: bearer, - cookie, - sentinel: None, - }; + sentinel = Some(SoraSentinel::new(value)); } - Ok(CredentialsPayload { - credentials_set: SoraCredentialSet::from_legacy_credentials(&credentials)?, - credentials, - }) + 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 985fbf6b87..4808e1ef03 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,14 +19,13 @@ pub struct SoraCredentialManager { impl SoraCredentialManager { - pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { + pub async fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { let holder = SoraCredentialHolder::new(); - match read_sora_credentials_from_disk(app_data_root) { + match read_sora_credentials_from_disk(app_data_root).await { Err(err) => warn!("Failed to read credentials from disk: {:?}", err), Ok(creds) => { - holder.set_legacy_credentials(&creds.credentials).expect("Failed to set credentials"); - holder.set_credentials(&creds.credentials_set).expect("Failed to set credentials"); + holder.set_credentials(&creds).expect("Failed to set credentials"); } } @@ -65,11 +64,10 @@ impl SoraCredentialManager { self.holder.get_credentials_required() } - pub fn reset_from_disk(&self) -> AnyhowResult { - let creds = read_sora_credentials_from_disk(&self.app_data_root)?; - self.holder.set_credentials(&creds.credentials_set)?; - self.holder.set_legacy_credentials(&creds.credentials)?; - Ok(creds.credentials) + pub async fn reset_from_disk(&self) -> AnyhowResult { + let creds = read_sora_credentials_from_disk(&self.app_data_root).await?; + self.holder.set_credentials(&creds)?; + Ok(creds) } /// Refresh the sentinel token from Sora's API diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs index 2d48e20af4..636c3cda04 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs @@ -13,6 +13,18 @@ pub struct SoraCredentialSet { } impl SoraCredentialSet { + pub fn initialize( + cookies: SoraCookies, + bearer: Option, + sentinel: Option + ) -> Self { + Self { + cookies, + jwt_bearer_token: bearer, + sora_sentinel: sentinel, + } + } + pub fn initialize_with_just_cookies(cookies: SoraCookies) -> Self { SoraCredentialSet { cookies,