diff --git a/crates/desktop/tauri-artcraft/Cargo.toml b/crates/desktop/tauri-artcraft/Cargo.toml index 781b17617a..6e3a1b4138 100644 --- a/crates/desktop/tauri-artcraft/Cargo.toml +++ b/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 } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs index 2038c13e56..d0ffa3f7c7 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs +++ b/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; \ No newline at end of file +pub mod sora_credential_manager; +mod sora_task_queue; \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs b/crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs new file mode 100644 index 0000000000..27ae17cd7b --- /dev/null +++ b/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 +} + +#[derive(Clone)] +pub struct SoraTaskQueue { + + queue: Arc>> +} diff --git a/crates/desktop/tauri-artcraft/src/threads/mod.rs b/crates/desktop/tauri-artcraft/src/threads/mod.rs index 5413233fa2..005f69d65c 100644 --- a/crates/desktop/tauri-artcraft/src/threads/mod.rs +++ b/crates/desktop/tauri-artcraft/src/threads/mod.rs @@ -1,2 +1,3 @@ pub mod discord_presence_thread; -pub mod sora_session_login_thread; \ No newline at end of file +pub mod sora_session_login_thread; +mod sora_task_polling_thread; \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs b/crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs new file mode 100644 index 0000000000..011f49ebd0 --- /dev/null +++ b/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; + } +} diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/image_remix_with_session_auto_renew.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/image_remix_with_session_auto_renew.rs index ee68c0e638..db28f678ce 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/image_remix_with_session_auto_renew.rs +++ b/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), SoraError> { let result = sora_image_gen_remix(SoraImageGenRemixRequest { prompt: request.prompt.clone(), // NB: Clone because used again diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_bytes_with_session_auto_renew.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_bytes_with_session_auto_renew.rs index a356a277be..2f8fba3feb 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_bytes_with_session_auto_renew.rs +++ b/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), SoraError> { diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_file_with_session_auto_renew.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_file_with_session_auto_renew.rs index 79b31bff0b..f4cf28934f 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/image_upload_from_file_with_session_auto_renew.rs +++ b/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> { } /// 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>( request: ImageUploadFromFileAutoRenewRequest<'_, P> ) -> Result<(SoraMediaUploadResponse, Option), SoraError> { diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/list_sora_task_status_with_session_auto_renew.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/list_sora_task_status_with_session_auto_renew.rs new file mode 100644 index 0000000000..b25a8219c3 --- /dev/null +++ b/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, + pub before: Option, + 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), 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))) +} diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs index db5d9f692d..e1c2ba6090 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs +++ b/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; \ No newline at end of file diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs index 79a9a71d89..50d5bf0b74 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs +++ b/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) } diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_status.rs b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_status.rs index f1012da827..54ba29e17c 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_status.rs +++ b/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, pub url: String, @@ -170,7 +182,7 @@ pub struct ModerationResult { pub results_by_frame_index: serde_json::Value, pub code: Option, 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 { +pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentials) -> Result { 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?; diff --git a/crates/lib/api_clients/openai_sora_client/src/sora_error.rs b/crates/lib/api_clients/openai_sora_client/src/sora_error.rs index 78fa2e56e7..86ecb17c28 100644 --- a/crates/lib/api_clients/openai_sora_client/src/sora_error.rs +++ b/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 for SoraError { Self::AnyhowError(err) } } + +impl From for SoraError { + fn from(err: serde_json::Error) -> Self { + Self::JsonError(err) + } +}