better credential behavior

crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs

@@ -1,8 +1,9 @@
-use base64::Engine;
-use base64::prelude::BASE64_STANDARD;
 use crate::state::app_dir::AppDataRoot;
 use crate::state::sora::read_sora_credentials_from_disk::read_sora_credentials_from_disk;
 use crate::state::sora::sora_credential_holder::SoraCredentialHolder;
+use crate::state::sora::sora_credential_manager::SoraCredentialManager;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
 use errors::AnyhowResult;
 use image::codecs::png::{CompressionType, FilterType, PngEncoder};
 use image::{DynamicImage, ImageReader};
@@ -10,6 +11,7 @@ use log::{error, info};
 use openai_sora_client::credentials::SoraCredentials;
 use openai_sora_client::image_gen::common::{ImageSize, NumImages};
 use openai_sora_client::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest};
+use openai_sora_client::sentinel_refresh::refresh_sentinel;
 use openai_sora_client::upload::upload_media_from_bytes::sora_media_upload_from_bytes;
 use openai_sora_client::upload::upload_media_from_file::SoraMediaUploadRequest;
 use serde_derive::Deserialize;
@@ -40,11 +42,11 @@ pub async fn sora_image_remix_command(
   _app: AppHandle,
   request: SoraImageRemixCommand,
   app_data_root: State<'_, AppDataRoot>,
-  sora_creds_holder: State<'_, SoraCredentialHolder>,
+  sora_creds_manager: State<'_, SoraCredentialManager>,
 ) -> Result<String, String> {
   info!("image_generation_command called; processing image...");
 
-  generate_image(request, &app_data_root, &sora_creds_holder)
+  generate_image(request, &app_data_root, &sora_creds_manager)
     .await
     .map_err(|err| {
       error!("error: {:?}", err);
@@ -57,16 +59,31 @@ pub async fn sora_image_remix_command(
 pub async fn generate_image(
   request: SoraImageRemixCommand,
   app_data_root: &AppDataRoot,
-  sora_creds_holder: &SoraCredentialHolder,
+  sora_creds_manager: &SoraCredentialManager,
 ) -> AnyhowResult<()> {
 
-  let sora_credentials = read_sora_credentials_from_disk(app_data_root)
+  // TODO(bt,2025-04-21): Read from in-memory cache instead, but allow for desktop replacement.
+  let mut sora_creds = read_sora_credentials_from_disk(app_data_root)
     .map_err(|err| {
       error!("Failed to read Sora credentials from disk: {:?}", err);
       err
     })?;
 
-  //let sora_credentials = sora_creds_holder.get_credentials()?;
+  let no_sentinel = !app_data_root.get_sora_sentinel_file_path().is_file()
+      || sora_creds.sentinel.is_none();
+
+  if no_sentinel {
+    info!("Refreshing credentials...");
+
+    sora_creds = sora_creds_manager.call_sentinel_refresh()
+        .await
+        .map_err(|err| {
+          error!("Failed to refresh: {:?}", err);
+          err
+        })?;
+
+    // TODO: Write to disk.
+  }
 
   let sora_media_tokens = vec![];
 
@@ -74,13 +91,36 @@ pub async fn generate_image(
   //  Note: This is incredibly inefficient. We should keep a local cache.
   //  Also, if they've already been uploaded to OpenAI, we shouldn't continue to re-upload.
 
-  let response = sora_image_gen_remix(SoraImageGenRemixRequest {
+  let mut response = sora_image_gen_remix(SoraImageGenRemixRequest {
     prompt: request.prompt.to_string(),
     num_images: NumImages::One,
     image_size: ImageSize::Square,
     sora_media_tokens: sora_media_tokens.clone(),
-    credentials: &sora_credentials,
-  }).await
+    credentials: &sora_creds,
+  }).await;
+
+  if let Err(err) = &response {
+    error!("Error in generating image: {:?}", err);
+
+    sora_creds = sora_creds_manager.call_sentinel_refresh()
+        .await
+        .map_err(|err| {
+          error!("Failed to refresh: {:?}", err);
+          err
+        })?;
+
+    info!("Retrying request...");
+
+    response = sora_image_gen_remix(SoraImageGenRemixRequest {
+      prompt: request.prompt.to_string(),
+      num_images: NumImages::One,
+      image_size: ImageSize::Square,
+      sora_media_tokens: sora_media_tokens.clone(),
+      credentials: &sora_creds,
+    }).await;
+  }
+
+  let response = response
       .map_err(|err| {
         error!("Failed to call Sora image generation: {:?}", err);
         err

crates/desktop/tauri-artcraft/src/lib.rs

@@ -11,6 +11,7 @@ use crate::commands::sora::sora_image_generation_command::sora_image_generation_
 use crate::commands::sora::sora_image_remix_command::sora_image_remix_command;
 use crate::state::app_config::AppConfig;
 use crate::state::sora::sora_credential_holder::SoraCredentialHolder;
+use crate::state::sora::sora_credential_manager::SoraCredentialManager;
 use crate::threads::sora_session_login_thread::sora_session_login_thread;
 
 use tauri_plugin_log::Target;
@@ -29,16 +30,9 @@ pub fn run() {
   let app_data_root = config.app_data_root.clone();
   let app_data_root2 = config.app_data_root.clone();
 
-  let sora_creds_holder = SoraCredentialHolder::new();
-  let sora_creds_holder2 = sora_creds_holder.clone();
-
   println!("Attempting to read existing credentials...");
-
-  if let Ok(creds) = read_sora_credentials_from_disk(&app_data_root) {
-    sora_creds_holder.set_credentials(creds).expect("there shouldn't be locking errors");
-  } else {
-    println!("No credentials found on disk.");
-  }
+  let sora_creds_manager = SoraCredentialManager::initialize_from_disk_infallible(&app_data_root);
+  let sora_creds_manager2 = sora_creds_manager.clone();
 
   println!("Initializing backend runtime...");
 
@@ -65,13 +59,13 @@ pub fn run() {
       //}
       let app = app.handle().clone();
 
-      tauri::async_runtime::spawn(sora_session_login_thread(app, app_data_root2, sora_creds_holder2));
+      tauri::async_runtime::spawn(sora_session_login_thread(app, app_data_root2, sora_creds_manager2));
 
       Ok(())
     })
     .manage(config)
     .manage(app_data_root)
-    .manage(sora_creds_holder)
+    .manage(sora_creds_manager)
     .invoke_handler(tauri::generate_handler![
       flip_image,
       open_sora_login_command,

crates/desktop/tauri-artcraft/src/state/sora/mod.rs

@@ -1,2 +1,3 @@
 pub mod read_sora_credentials_from_disk;
-pub mod sora_credential_holder;
+pub mod sora_credential_holder;
+pub mod sora_credential_manager;

crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs

@@ -1,4 +1,5 @@
-use std::sync::{Arc, RwLock};
+use std::ops::Deref;
+use std::sync::{Arc, RwLock, RwLockReadGuard};
 use anyhow::anyhow;
 use errors::AnyhowResult;
 use openai_sora_client::credentials::SoraCredentials;
@@ -16,11 +17,11 @@ impl SoraCredentialHolder {
     }
   }
 
-  pub fn set_credentials(&self, credentials: SoraCredentials) -> AnyhowResult<()> {
+  pub fn set_credentials(&self, credentials: &SoraCredentials) -> AnyhowResult<()> {
     match self.credentials.write() {
       Err(err) => Err(anyhow!("Failed to acquire write lock: {:?}", err)),
       Ok(mut creds) => {
-        *creds = Some(credentials);
+        *creds = Some(credentials.clone());
         Ok(())
       }
     }
@@ -44,4 +45,14 @@ impl SoraCredentialHolder {
       }
     }
   }
+
+  pub fn get_credentials_required(&self) -> AnyhowResult<SoraCredentials> {
+    match self.credentials.read() {
+      Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)),
+      Ok(creds) => match creds.deref() {
+        None => Err(anyhow!("Credentials not set")),
+        Some(creds) => Ok(creds.clone()),
+      }
+    }
+  }
 }

crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs

@@ -0,0 +1,83 @@
+use crate::state::app_dir::AppDataRoot;
+use crate::state::sora::read_sora_credentials_from_disk::read_sora_credentials_from_disk;
+use crate::state::sora::sora_credential_holder::SoraCredentialHolder;
+use errors::AnyhowResult;
+use log::{error, warn};
+use openai_sora_client::credentials::SoraCredentials;
+use openai_sora_client::sentinel_refresh::refresh_sentinel;
+use std::fs::OpenOptions;
+use std::io::Write;
+
+#[derive(Clone)]
+pub struct SoraCredentialManager {
+  holder: SoraCredentialHolder,
+  app_data_root: AppDataRoot,
+}
+
+impl SoraCredentialManager {
+
+  pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self {
+    let holder = SoraCredentialHolder::new();
+
+    match read_sora_credentials_from_disk(app_data_root) {
+      Err(err) => warn!("Failed to read credentials from disk: {:?}", err),
+      Ok(creds) => {
+        holder.set_credentials(&creds).expect("Failed to set credentials");
+      }
+    }
+
+    Self {
+      holder,
+      app_data_root: app_data_root.clone(),
+    }
+  }
+
+  pub fn set_credentials(&self, creds: &SoraCredentials) -> AnyhowResult<()> {
+    self.holder.set_credentials(creds)
+  }
+
+  pub fn clear_credentials(&self) -> AnyhowResult<()> {
+    self.holder.clear_credentials()
+  }
+
+  pub fn get_credentials(&self) -> AnyhowResult<Option<SoraCredentials>> {
+    self.holder.get_credentials()
+  }
+
+  pub fn get_credentials_required(&self) -> AnyhowResult<SoraCredentials> {
+    self.holder.get_credentials_required()
+  }
+
+  pub fn reset_from_disk(&self) -> AnyhowResult<SoraCredentials> {
+    let creds = read_sora_credentials_from_disk(&self.app_data_root)?;
+    self.holder.set_credentials(&creds)?;
+    Ok(creds)
+  }
+
+  pub async fn call_sentinel_refresh(&self) -> AnyhowResult<SoraCredentials> {
+    // NB(bt,2025-04-21): Technically we don't need credentials to get a sentinel.
+    let mut creds = self.holder.get_credentials_required()?;
+
+    let sentinel = refresh_sentinel()
+        .await
+        .map_err(|err| {
+          error!("Failed to refresh: {:?}", err);
+          err
+        })?;
+
+    creds.sentinel = Some(sentinel.clone());
+
+    self.holder.set_credentials(&creds)?;
+
+    let mut file = OpenOptions::new()
+        .create(true)
+        .write(true)
+        .truncate(true)
+        .open(self.app_data_root.get_sora_sentinel_file_path())?;
+
+    file.write_all(sentinel.as_bytes())?;
+    file.flush()?;
+
+    Ok(creds)
+  }
+}

crates/desktop/tauri-artcraft/src/threads/sora_session_login_thread.rs

@@ -8,6 +8,7 @@ use tauri::{AppHandle, Manager, Webview};
 use errors::AnyhowResult;
 use crate::state::app_dir::AppDataRoot;
 use crate::state::sora::sora_credential_holder::SoraCredentialHolder;
+use crate::state::sora::sora_credential_manager::SoraCredentialManager;
 use crate::utils::sora_webview_cookies::get_all_sora_cookies_as_string;
 
 pub const LOGIN_WINDOW_NAME: &str = "login_window";
@@ -27,7 +28,7 @@ pub static SORA_ROOT_URL: Lazy<Url> = Lazy::new(|| {
 pub async fn sora_session_login_thread(
   app: AppHandle,
   app_data_root: AppDataRoot,
-  sora_creds_holder: SoraCredentialHolder
+  sora_creds_manager: SoraCredentialManager
 ) -> ! {
   loop {
     for (window_name, webview) in app.webviews() {