auto-healing task polling, beginnings of a task polling thread

crates/desktop/tauri-artcraft/Cargo.toml

@@ -34,8 +34,6 @@ subprocess_common = { workspace = true }
 tokens = { workspace = true }
 
 # External
-#openssl = { version = "0.10", features = ["vendored"] }
-# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
 anyhow = { workspace = true }
 base64 = "0.22.1"
 bytes = "1.10"
@@ -48,6 +46,7 @@ image = "0.25.5"
 log = { workspace = true }
 once_cell = { workspace = true }
 rand = { workspace = true }
+# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
 reqwest = { workspace = true, features = ["multipart", "stream"], default-features = false }
 serde = { workspace = true }
 serde_derive = { workspace = true }

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

@@ -1,4 +1,5 @@
 pub mod read_sora_credentials_from_disk;
 pub mod read_sora_credentials_from_disk_with_upgrade;
 pub mod sora_credential_holder;
-pub mod sora_credential_manager;
+pub mod sora_credential_manager;
+mod sora_task_queue;

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

@@ -0,0 +1,14 @@
+use chrono::{Date, DateTime, Utc};
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+#[derive(Clone)]
+pub struct TaskDetails {
+  pub enqueue_time: DateTime<Utc>
+}
+
+#[derive(Clone)]
+pub struct SoraTaskQueue {
+
+  queue: Arc<RwLock<HashMap<String, TaskDetails>>>
+}

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

@@ -1,2 +1,3 @@
 pub mod discord_presence_thread;
-pub mod sora_session_login_thread;
+pub mod sora_session_login_thread;
+mod sora_task_polling_thread;

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

@@ -0,0 +1,44 @@
+use crate::state::app_dir::AppDataRoot;
+use crate::state::sora::sora_credential_manager::SoraCredentialManager;
+use errors::AnyhowResult;
+use log::{error, info};
+use openai_sora_client::recipes::list_sora_task_status_with_session_auto_renew::{list_sora_task_status_with_session_auto_renew, StatusRequestArgs};
+use tauri::AppHandle;
+
+pub async fn sora_session_login_thread(
+  app_data_root: AppDataRoot,
+  sora_creds_manager: SoraCredentialManager
+) -> ! {
+  loop {
+    if let Err(err) = polling_loop(&sora_creds_manager).await {
+      error!("An error occurred: {:?}", err);
+    }
+    tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
+  }
+}
+
+async fn polling_loop(sora_creds_manager: &SoraCredentialManager) -> AnyhowResult<()> {
+  loop {
+    let creds = sora_creds_manager.get_credentials_required()?;
+
+    let (response, maybe_new_creds) = list_sora_task_status_with_session_auto_renew(StatusRequestArgs {
+      limit: None,
+      before: None,
+      credentials: creds,
+    }).await?;
+
+    if let Some(new_creds) = maybe_new_creds {
+      info!("Saving new credentials.");
+      sora_creds_manager.set_credentials(&new_creds)?;
+    }
+
+    // TODO: The cursoring logic likely needs to improve.
+    for task in response.task_responses {
+      info!("Task ID: {}", task.id);
+      info!("Task Status: {:?}", task.status);
+      // Handle the task status as needed
+    }
+
+    tokio::time::sleep(std::time::Duration::from_millis(10_000)).await;
+  }
+}

crates/lib/api_clients/openai_sora_client/src/recipes/image_remix_with_session_auto_renew.rs

@@ -21,6 +21,7 @@ pub struct ImageRemixAutoRenewRequest<'a> {
 }
 
 /// Image request with retry and session auto-renewal.
+/// If a new sora credential is returned, replace the old one with the new one.
 pub async fn image_remix_with_session_auto_renew(request: ImageRemixAutoRenewRequest<'_>) -> Result<(SoraImageGenResponse, Option<SoraCredentialSet>), SoraError> {
   let result = sora_image_gen_remix(SoraImageGenRemixRequest {
     prompt: request.prompt.clone(), // NB: Clone because used again

crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_bytes_with_session_auto_renew.rs

@@ -23,6 +23,7 @@ pub struct ImageUploadFromBytesAutoRenewRequest<'a> {
 }
 
 /// Image upload with retry and session auto-renewal.
+/// If a new sora credential is returned, replace the old one with the new one.
 pub async fn image_upload_from_bytes_with_session_auto_renew(
   request: ImageUploadFromBytesAutoRenewRequest<'_>
 ) -> Result<(SoraMediaUploadResponse, Option<SoraCredentialSet>), SoraError> {

crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_file_with_session_auto_renew.rs

@@ -13,6 +13,7 @@ pub struct ImageUploadFromFileAutoRenewRequest<'a, P: AsRef<Path>> {
 }
 
 /// Image upload with retry and session auto-renewal.
+/// If a new sora credential is returned, replace the old one with the new one.
 pub async fn image_upload_from_file_with_session_auto_renew<P: AsRef<Path>>(
   request: ImageUploadFromFileAutoRenewRequest<'_, P>
 ) -> Result<(SoraMediaUploadResponse, Option<SoraCredentialSet>), SoraError> {

crates/lib/api_clients/openai_sora_client/src/recipes/list_sora_task_status_with_session_auto_renew.rs

@@ -0,0 +1,62 @@
+use crate::creds::credential_migration::CredentialMigrationRef;
+use crate::creds::sora_credential_set::SoraCredentialSet;
+use crate::recipes::image_upload_from_file_with_session_auto_renew::ImageUploadFromFileAutoRenewRequest;
+use crate::recipes::maybe_refresh_credentials_on_sora_error::maybe_refresh_credentials_on_sora_error;
+use crate::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
+use crate::requests::image_gen::image_gen_status::{get_image_gen_status, StatusRequest, TaskId, TaskResponse, VideoGenStatusResponse};
+use crate::requests::upload::upload_media_from_file::sora_media_upload_from_file;
+use crate::requests::upload::upload_media_http_request::SoraMediaUploadResponse;
+use crate::sora_error::SoraError;
+use anyhow::Error;
+use log::{info, warn};
+use std::path::Path;
+
+pub struct StatusRequestArgs {
+  pub limit: Option<u32>,
+  pub before: Option<TaskId>,
+  pub credentials: SoraCredentialSet,
+}
+
+/// Check Sora task statuses with session auto-renewal.
+/// If a new sora credential is returned, replace the old one with the new one.
+pub async fn list_sora_task_status_with_session_auto_renew(
+  args: StatusRequestArgs,
+) -> Result<(VideoGenStatusResponse, Option<SoraCredentialSet>), SoraError> {
+
+  let request = StatusRequest {
+    limit: args.limit,
+    before: args.before.map(|t| t.0),
+  };
+
+  let mut new_creds = None;
+
+  if args.credentials.jwt_bearer_token.is_none() {
+    info!("JWT not set. Upgrading credentials...");
+    let mut updated_creds = args.credentials.clone();
+    if let Err(err) = maybe_upgrade_or_renew_session(&mut updated_creds).await {
+      warn!("Failed to upgrade or renew session: {:?}", err);
+    }
+    new_creds = Some(updated_creds);
+  }
+
+  let legacy_creds = new_creds.as_ref()
+      .unwrap_or_else(|| &args.credentials)
+      .to_legacy_credentials()?;
+
+  let result = get_image_gen_status(&request, &legacy_creds).await;
+
+  let err = match result {
+    Ok(response) => return Ok((response, new_creds)),
+    Err(err) => err,
+  };
+
+  warn!("Task status polling failed: {:?}", err);
+
+  let new_creds = maybe_refresh_credentials_on_sora_error(&args.credentials, err).await?;
+
+  info!("Retrying image upload with new credentials...");
+
+  let result = get_image_gen_status(&request, &legacy_creds).await?;
+
+  Ok((result, Some(new_creds)))
+}

crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs

@@ -1,6 +1,7 @@
 pub mod image_remix_with_session_auto_renew;
 pub mod image_upload_from_bytes_with_session_auto_renew;
 pub mod image_upload_from_file_with_session_auto_renew;
+pub mod list_sora_task_status_with_session_auto_renew;
 pub mod maybe_refresh_credentials_on_sora_error;
 pub mod maybe_upgrade_or_renew_session;
 pub mod wait_for_image_gen_status;

crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs

@@ -47,5 +47,5 @@ pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCrede
 }
 
 fn find_task_response_by_id(status_response: &VideoGenStatusResponse, task_id: String) -> Option<&TaskResponse> {
-  status_response.task_responses.iter().find(|task| task.id == task_id)
+  status_response.task_responses.iter().find(|task| task.id.0 == task_id)
 }

crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_status.rs

@@ -1,15 +1,27 @@
-use std::{fs::File, io::Write, path::Path};
-
 use crate::credentials::SoraCredentials;
+use crate::sora_error::SoraError;
+use anyhow::anyhow;
 use errors::AnyhowResult;
 use reqwest::Url;
-use serde_derive::Deserialize;
+use serde_derive::{Deserialize, Serialize};
+use std::fmt::{Debug, Display};
+use std::{fs::File, io::Write, path::Path};
 
 const SORA_STATUS_URL: &str = "https://sora.com/backend/video_gen";
 
 /// This user agent is tied to the sentinel generation. If we need to change it, we may need to change sentinel generation too.
 const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36";
 
+/// A strongly typed task ID for Sora tasks.
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize, Deserialize)]
+pub struct TaskId(pub String);
+
+impl Display for TaskId {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    std::fmt::Display::fmt(&self.0, f)
+  }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
 pub enum TaskStatus {
   Queued,
@@ -42,7 +54,7 @@ pub struct VideoGenStatusResponse {
 #[derive(Deserialize, Clone, Debug)]
 #[serde(rename_all = "snake_case")]
 pub struct TaskResponse {
-  pub id: String,
+  pub id: TaskId,
   pub user: String,
   pub created_at: String,
   pub status: String,
@@ -79,7 +91,7 @@ pub struct TaskResponse {
 #[serde(rename_all = "snake_case")]
 pub struct Generation {
   pub id: String,
-  pub task_id: String,
+  pub task_id: TaskId,
   pub created_at: String,
   pub deleted_at: Option<String>,
   pub url: String,
@@ -170,7 +182,7 @@ pub struct ModerationResult {
   pub results_by_frame_index: serde_json::Value,
   pub code: Option<String>,
   pub is_output_rejection: bool,
-  pub task_id: String,
+  pub task_id: TaskId,
 }
 
 pub struct StatusRequest {
@@ -186,10 +198,11 @@ impl StatusRequest {
 }
 
 /// Gets the status of image generation tasks from Sora API
-pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentials) -> AnyhowResult<VideoGenStatusResponse> {
+pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentials) -> Result<VideoGenStatusResponse, SoraError> {
   let client = reqwest::Client::new();
 
-  let mut url = reqwest::Url::parse(SORA_STATUS_URL)?;
+  let mut url = reqwest::Url::parse(SORA_STATUS_URL)
+      .map_err(|err| anyhow!("error parsing URL: {:?}", err))?;
 
   // Add query parameters
   if let Some(limit) = status_request.limit {
@@ -200,7 +213,11 @@ pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &
     url.query_pairs_mut().append_pair("before", &before);
   }
 
-  let http_request = client.get(url).header("User-Agent", USER_AGENT).header("Cookie", &credentials.cookie).header("Authorization", credentials.authorization_header_value()).header("Content-Type", "application/json");
+  let http_request = client.get(url)
+      .header("User-Agent", USER_AGENT)
+      .header("Cookie", &credentials.cookie)
+      .header("Authorization", credentials.authorization_header_value())
+      .header("Content-Type", "application/json");
 
   let http_request = http_request.build()?;
 
@@ -212,7 +229,8 @@ pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &
   println!("response_body: {}", response_body);
 
   if status != reqwest::StatusCode::OK {
-    return Err(anyhow::anyhow!("The status request failed; status = {:?}, message = {:?}", status, response_body));
+    // TODO: Handle non-200 status codes, esp. JWT errors and credential errors
+    return Err(SoraError::AnyhowError(anyhow::anyhow!("The status request failed; status = {:?}, message = {:?}", status, response_body)));
   }
 
   let response: VideoGenStatusResponse = serde_json::from_str(response_body)?;
@@ -285,7 +303,7 @@ mod tests {
 
     let task_id = response.task_responses[0].id.clone();
     let task_response = wait_for_image_gen_status(
-      &task_id,
+      &task_id.0,
       &creds,
       Some(10)
     ).await?;

crates/lib/api_clients/openai_sora_client/src/sora_error.rs

@@ -5,12 +5,14 @@ use std::io;
 /// The single error type this crate should surface to all callers.
 #[derive(Debug)]
 pub enum SoraError {
-  /// We haven't received a bearer token yet (this is our application error)
+  /// We haven't received a bearer token yet
+  /// This is our own internal application state error, not something Sora returns.
   NoBearerTokenAvailable,
 
   /// We're sending too many tasks to Sora.
   TooManyConcurrentTasks,
 
+  // TODO: Try to make concrete error types for these.
   /// Unauthorized, cookie and/or bearer token expired. We'll need to ask for a refreshed login.
   ///
   /// Example message, e.g. from the upload endpoint:
@@ -27,7 +29,7 @@ pub enum SoraError {
   /// Another error occurred.
   OtherBadStatus(anyhow::Error),
 
-  /// Reqwest Error
+  /// Reqwest Error, eg. connection failures.
   ReqwestError(reqwest::Error),
 
   /// std::io Error that arises from our end, eg. reading from the filesystem.
@@ -35,6 +37,9 @@ pub enum SoraError {
 
   /// anyhow::Error arises from our end
   AnyhowError(anyhow::Error),
+
+  /// serde_json::Error, likely from JSON deserialization schema mismatch.
+  JsonError(serde_json::Error),
 }
 
 impl Error for SoraError {}
@@ -63,6 +68,9 @@ impl Display for SoraError {
       Self::AnyhowError(err) => {
         write!(f, "Anyhow error: {}", err)
       }
+      Self::JsonError(err) => {
+        write!(f, "serde_json error: {}", err)
+      }
     }
   }
 }
@@ -84,3 +92,9 @@ impl From<anyhow::Error> for SoraError {
     Self::AnyhowError(err)
   }
 }
+
+impl From<serde_json::Error> for SoraError {
+  fn from(err: serde_json::Error) -> Self {
+    Self::JsonError(err)
+  }
+}