From ee4de4126d97dcb4afa87c53cb89ff245e8aa1be Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 12 May 2025 12:26:41 -0400 Subject: [PATCH] timeouts for uploads / remix (#602) --- .../sora/sora_image_generation_command.rs | 13 ++++++++++- .../commands/sora/sora_image_remix_command.rs | 21 ++++++++++++------ .../image_remix_with_session_auto_renew.rs | 5 +++++ ...load_from_bytes_with_session_auto_renew.rs | 6 +++++ ...pload_from_file_with_session_auto_renew.rs | 7 ++++++ .../image_gen/image_gen_http_request.rs | 22 ++++++++++++++----- .../image_gen/sora_image_gen_remix.rs | 7 ++++-- .../image_gen/sora_image_gen_simple.rs | 10 +++++---- .../upload/upload_media_from_bytes.rs | 10 +++++++-- .../requests/upload/upload_media_from_file.rs | 10 +++++++-- .../upload/upload_media_http_request.rs | 8 ++++++- 11 files changed, 94 insertions(+), 25 deletions(-) diff --git a/crates/desktop/artcraft/src/commands/sora/sora_image_generation_command.rs b/crates/desktop/artcraft/src/commands/sora/sora_image_generation_command.rs index 9f42a82d41..1eeea144bc 100644 --- a/crates/desktop/artcraft/src/commands/sora/sora_image_generation_command.rs +++ b/crates/desktop/artcraft/src/commands/sora/sora_image_generation_command.rs @@ -14,8 +14,13 @@ use openai_sora_client::requests::image_gen::sora_image_gen_remix::{sora_image_g use openai_sora_client::requests::upload::upload_media_from_bytes::sora_media_upload_from_bytes; use std::fs::read_to_string; use std::io::Cursor; +use std::time::Duration; use tauri::{AppHandle, Manager, State}; +const SORA_IMAGE_UPLOAD_TIMEOUT: Duration = Duration::from_millis(1000 * 30); // 30 seconds + +const SORA_IMAGE_REMIX_TIMEOUT: Duration = Duration::from_millis(1000 * 30); // 30 seconds + #[tauri::command] pub async fn sora_image_generation_command( prompt: &str, @@ -55,8 +60,13 @@ pub async fn generate_image( let image_bytes = BASE64_STANDARD.decode(image)?; let filename = "image.png".to_string(); + - let response = sora_media_upload_from_bytes(image_bytes, filename, CredentialMigrationRef::New(&creds)) + let response = sora_media_upload_from_bytes( + image_bytes, + filename, + CredentialMigrationRef::New(&creds), + Some(SORA_IMAGE_UPLOAD_TIMEOUT)) .await .map_err(|err| { error!("Failed to upload image to Sora: {:?}", err); @@ -72,6 +82,7 @@ pub async fn generate_image( image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), credentials: CredentialMigrationRef::New(&creds), + request_timeout: Some(SORA_IMAGE_REMIX_TIMEOUT), }).await .map_err(|err| { error!("Failed to call Sora image generation: {:?}", err); diff --git a/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs b/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs index 8120a4e42e..b7ef12c517 100644 --- a/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs +++ b/crates/desktop/artcraft/src/commands/sora/sora_image_remix_command.rs @@ -25,17 +25,22 @@ 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; use openai_sora_client::requests::upload::upload_media_from_file::sora_media_upload_from_file; +use openai_sora_client::sora_error::SoraError; use serde_derive::{Deserialize, Serialize}; use std::fs::read_to_string; use std::io::Cursor; +use std::time::Duration; +use storyteller_client::api_error::ApiError; +use storyteller_client::api_error::ApiError::InternalServerError; use storyteller_client::media_files::get_media_file::get_media_file; use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Emitter, Manager, State}; -use openai_sora_client::sora_error::SoraError; -use storyteller_client::api_error::ApiError; -use storyteller_client::api_error::ApiError::InternalServerError; use tokens::tokens::media_files::MediaFileToken; +const SORA_IMAGE_UPLOAD_TIMEOUT: Duration = Duration::from_millis(1000 * 30); // 30 seconds + +const SORA_IMAGE_REMIX_TIMEOUT: Duration = Duration::from_millis(1000 * 30); // 30 seconds + #[derive(Deserialize)] pub struct SoraImageRemixCommand { /// Image media file; the engine or canvas snapshot (screenshot). @@ -96,23 +101,23 @@ pub async fn sora_image_remix_command( let mut status = CommandErrorStatus::ServerError; let mut error_type = SoraImageRemixErrorType::ServerError; - let mut error_message = "A server error occurred."; + let mut error_message = "A server error occurred. Please try again. If it continues, please tell our staff about the problem."; match (err) { InnerError::SoraError(SoraError::TooManyConcurrentTasks) => { error_type = SoraImageRemixErrorType::TooManyConcurrentTasks; status = CommandErrorStatus::ServerError; - error_message = "You already have work in progress. Please wait."; + error_message = "You already have work in progress. Please wait for it to finish."; }, InnerError::SoraError(SoraError::SoraUsernameNotYetCreated) => { error_type = SoraImageRemixErrorType::SoraUsernameNotYetCreated; status = CommandErrorStatus::BadRequest; - error_message = "You need to create a username on Sora.com."; + error_message = "You need to create a username on Sora.com to continue."; }, InnerError::SoraError(SoraError::BadGateway(_) | SoraError::CloudFlareTimeout(_)) => { error_type = SoraImageRemixErrorType::SoraIsHavingProblems; status = CommandErrorStatus::ServerError; - error_message = "Sora is having problems. Please wait a moment before retrying."; + error_message = "Sora is having problems. Please wait a moment and then retry."; } _ => {}, } @@ -206,6 +211,7 @@ pub async fn generate_image( image_upload_from_file_with_session_auto_renew(ImageUploadFromFileAutoRenewRequest { file_path, credentials: &creds, + request_timeout: Some(SORA_IMAGE_UPLOAD_TIMEOUT), }).await?; if let Some(new_creds) = maybe_new_credentials { @@ -230,6 +236,7 @@ pub async fn generate_image( image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), credentials: &creds, + request_timeout: Some(SORA_IMAGE_REMIX_TIMEOUT), }).await?; if let Some(new_creds) = maybe_new_creds { 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 080906af45..eb17d90d4c 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 @@ -1,3 +1,4 @@ +use std::time::Duration; use anyhow::anyhow; use log::{info, warn}; use errors::AnyhowResult; @@ -18,6 +19,8 @@ pub struct ImageRemixAutoRenewRequest<'a> { pub image_size: ImageSize, pub sora_media_tokens: Vec, pub credentials: &'a SoraCredentialSet, + /// This function can try several different requests. This is the timeout for *individual* requests. + pub request_timeout: Option, } /// Image request with retry and session auto-renewal. @@ -29,6 +32,7 @@ pub async fn image_remix_with_session_auto_renew(request: ImageRemixAutoRenewReq image_size: request.image_size, sora_media_tokens: request.sora_media_tokens.clone(), // NB: Clone because used again credentials: CredentialMigrationRef::New(request.credentials), + request_timeout: request.request_timeout, }).await; let err = match result { @@ -126,6 +130,7 @@ pub async fn image_remix_with_session_auto_renew(request: ImageRemixAutoRenewReq image_size: request.image_size, sora_media_tokens: request.sora_media_tokens, credentials: CredentialMigrationRef::New(&new_creds), + request_timeout: request.request_timeout, }).await; match result { 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 2f8fba3feb..1b06065637 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 @@ -14,12 +14,16 @@ use crate::sora_error::SoraError; use anyhow::anyhow; use errors::AnyhowResult; use log::{error, info, warn}; +use std::time::Duration; pub struct ImageUploadFromBytesAutoRenewRequest<'a> { pub file_bytes: Vec, pub filename: String, pub mime_type: &'a str, pub credentials: &'a SoraCredentialSet, + + /// This function can try to upload several times. This is the timeout for *individual* requests. + pub request_timeout: Option, } /// Image upload with retry and session auto-renewal. @@ -32,6 +36,7 @@ pub async fn image_upload_from_bytes_with_session_auto_renew( request.file_bytes.clone(), // FIXME(bt): This is horrible, but the client needs to take ownership. :( request.filename.clone(), // FIXME: Same CredentialMigrationRef::New(request.credentials), + request.request_timeout, ).await; let err = match result { @@ -51,6 +56,7 @@ pub async fn image_upload_from_bytes_with_session_auto_renew( request.file_bytes, request.filename, CredentialMigrationRef::New(&new_creds), + request.request_timeout, ).await?; Ok((result, Some(new_creds))) 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 f4cf28934f..695400f7db 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 @@ -6,10 +6,14 @@ use crate::requests::upload::upload_media_http_request::{upload_media_http_reque use crate::sora_error::SoraError; use log::{info, warn}; use std::path::Path; +use std::time::Duration; pub struct ImageUploadFromFileAutoRenewRequest<'a, P: AsRef> { pub file_path: P, pub credentials: &'a SoraCredentialSet, + + /// This function can try to upload several times. This is the timeout for *individual* requests. + pub request_timeout: Option, } /// Image upload with retry and session auto-renewal. @@ -21,6 +25,8 @@ pub async fn image_upload_from_file_with_session_auto_renew>( let result = sora_media_upload_from_file( request.file_path.as_ref().clone(), // FIXME(bt): This is horrible, but the client needs to take ownership. :( CredentialMigrationRef::New(request.credentials), + /// This function can try to upload several times. This is the timeout for *individual* requests. + request.request_timeout, ).await; let err = match result { @@ -37,6 +43,7 @@ pub async fn image_upload_from_file_with_session_auto_renew>( let result = sora_media_upload_from_file( request.file_path, CredentialMigrationRef::New(&new_creds), + request.request_timeout, ).await?; Ok((result, Some(new_creds))) diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_http_request.rs b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_http_request.rs index c446840548..b09094f42f 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_http_request.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/image_gen_http_request.rs @@ -1,8 +1,9 @@ -use log::warn; use crate::creds::credential_migration::CredentialMigrationRef; -use serde_derive::{Deserialize, Serialize}; -use thiserror::Error; use crate::requests::image_gen::image_gen_status::TaskId; +use log::warn; +use serde_derive::{Deserialize, Serialize}; +use std::time::Duration; +use thiserror::Error; const SORA_IMAGE_GEN_URL: &str = "https://sora.com/backend/video_gen"; @@ -152,7 +153,11 @@ impl std::fmt::Display for RawSoraErrorInner { } /// Don't expose the internal request implementation as there are only a few "correct" ways to call the API. -pub (crate) async fn image_gen_http_request(sora_request: RawSoraImageGenRequest, credentials: CredentialMigrationRef<'_>) -> Result { +pub (crate) async fn image_gen_http_request( + sora_request: RawSoraImageGenRequest, + credentials: CredentialMigrationRef<'_>, + request_timeout: Option, +) -> Result { let client = reqwest::Client::new(); let mut cookie; @@ -181,12 +186,17 @@ pub (crate) async fn image_gen_http_request(sora_request: RawSoraImageGenRequest } } - let http_request = client.post(SORA_IMAGE_GEN_URL) + let mut http_request = client.post(SORA_IMAGE_GEN_URL) .header("User-Agent", USER_AGENT) .header("Cookie", &cookie) .header("Authorization", &authorization_header) .header("Content-Type", "application/json") .header("OpenAI-Sentinel-Token", &sentinel); + + + if let Some(timeout) = request_timeout { + http_request = http_request.timeout(timeout); + } let http_request = http_request.json(&sora_request).build() .map_err(|e| SoraImageGenError::NetworkError(e.to_string()))?; @@ -241,8 +251,8 @@ pub (crate) async fn image_gen_http_request(sora_request: RawSoraImageGenRequest #[cfg(test)] mod tests { - use errors::AnyhowResult; use crate::requests::image_gen::image_gen_http_request::RawSoraResponse; + use errors::AnyhowResult; #[test] fn deserialize_task_id() -> AnyhowResult<()> { diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_remix.rs b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_remix.rs index 51422bb797..819aa49403 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_remix.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_remix.rs @@ -1,7 +1,8 @@ -use log::info; use crate::creds::credential_migration::CredentialMigrationRef; use crate::requests::image_gen::common::{ImageSize, NumImages, SoraImageGenResponse}; use crate::requests::image_gen::image_gen_http_request::{image_gen_http_request, InpaintItem, InpaintItemType, OperationType, RawSoraImageGenRequest, SoraImageGenError, VideoGenType}; +use log::info; +use std::time::Duration; pub struct SoraImageGenRemixRequest<'a> { pub prompt: String, @@ -9,6 +10,7 @@ pub struct SoraImageGenRemixRequest<'a> { pub image_size: ImageSize, pub sora_media_tokens: Vec, pub credentials: CredentialMigrationRef<'a>, + pub request_timeout: Option, } /// The "remix" commands let you supply additional images as context. @@ -39,7 +41,7 @@ pub async fn sora_image_gen_remix(request: SoraImageGenRemixRequest<'_>) -> Resu // TODO: Replace this with debug level logging in the future. info!("Sending remix request: {:?}", args); - let result = image_gen_http_request(args, request.credentials).await?; + let result = image_gen_http_request(args, request.credentials, request.request_timeout).await?; Ok(SoraImageGenResponse { task_id: result.id, @@ -80,6 +82,7 @@ mod tests { image_size: ImageSize::Square, sora_media_tokens: vec!["media_01jqyhrz4detyvtzwp2p4j63ad".to_string()], credentials: CredentialMigrationRef::Legacy(&creds), + request_timeout: None, }).await?; println!("task_id: {}", response.task_id); diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_simple.rs b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_simple.rs index 2fc3d34af6..93e247f230 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_simple.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/sora_image_gen_simple.rs @@ -1,14 +1,15 @@ -use crate::credentials::SoraCredentials; +use crate::creds::credential_migration::CredentialMigrationRef; use crate::requests::image_gen::common::{ImageSize, NumImages, SoraImageGenResponse}; use crate::requests::image_gen::image_gen_http_request::{image_gen_http_request, OperationType, RawSoraImageGenRequest, VideoGenType}; use errors::AnyhowResult; -use crate::creds::credential_migration::CredentialMigrationRef; +use std::time::Duration; pub struct SoraImageGenSimpleRequest<'a> { pub prompt: String, pub num_images: NumImages, pub image_size: ImageSize, pub credentials: CredentialMigrationRef<'a>, + pub request_timeout: Option, } pub async fn sora_image_gen_simple(args: SoraImageGenSimpleRequest<'_>) -> AnyhowResult { @@ -24,7 +25,7 @@ pub async fn sora_image_gen_simple(args: SoraImageGenSimpleRequest<'_>) -> Anyho }; // TODO: Error handling. - let result = image_gen_http_request(sora_request, args.credentials).await?; + let result = image_gen_http_request(sora_request, args.credentials, args.request_timeout).await?; Ok(SoraImageGenResponse { task_id: result.id, @@ -34,12 +35,12 @@ pub async fn sora_image_gen_simple(args: SoraImageGenSimpleRequest<'_>) -> Anyho #[cfg(test)] mod tests { use crate::credentials::SoraCredentials; + use crate::creds::credential_migration::CredentialMigrationRef; use crate::requests::image_gen::common::{ImageSize, NumImages}; use crate::requests::image_gen::sora_image_gen_simple::{sora_image_gen_simple, SoraImageGenSimpleRequest}; use errors::AnyhowResult; use std::fs::read_to_string; use testing::test_file_path::test_file_path; - use crate::creds::credential_migration::CredentialMigrationRef; #[ignore] // You can manually run "ignore" tests in the IDE, but they won't run in CI. #[tokio::test] @@ -64,6 +65,7 @@ mod tests { num_images: NumImages::One, image_size: ImageSize::Square, credentials: CredentialMigrationRef::Legacy(&creds), + request_timeout: None, }).await?; println!("task_id: {}", response.task_id); diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_bytes.rs b/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_bytes.rs index 1902b9ada9..686bca9261 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_bytes.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_bytes.rs @@ -2,10 +2,16 @@ use crate::creds::credential_migration::CredentialMigrationRef; use crate::requests::upload::upload_media_http_request::{upload_media_http_request, SoraMediaUploadResponse}; use crate::sora_error::SoraError; use std::path::PathBuf; +use std::time::Duration; /// Upload bytes. /// The underlying reqwest lib needs to own the bytes, so we can't pass as a reference. -pub async fn sora_media_upload_from_bytes(bytes: Vec, file_name: String, creds: CredentialMigrationRef<'_>) -> Result { +pub async fn sora_media_upload_from_bytes( + bytes: Vec, + file_name: String, + creds: CredentialMigrationRef<'_>, + maybe_timeout: Option, +) -> Result { let file_path = PathBuf::from(&file_name); // TODO: Read file magic bytes first, then fall back to this. @@ -20,5 +26,5 @@ pub async fn sora_media_upload_from_bytes(bytes: Vec, file_name: String, cre _ => "application/octet-stream", }; - Ok(upload_media_http_request(bytes, file_name, mime_type, creds).await?) + Ok(upload_media_http_request(bytes, file_name, mime_type, creds, maybe_timeout).await?) } diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_file.rs b/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_file.rs index dc023285c1..e7fb48cf0b 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_file.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_from_file.rs @@ -2,6 +2,7 @@ use crate::creds::credential_migration::CredentialMigrationRef; use crate::requests::upload::upload_media_http_request::{upload_media_http_request, SoraMediaUploadResponse}; use crate::sora_error::SoraError; use std::path::Path; +use std::time::Duration; use tokio::fs::File; use tokio::io::AsyncReadExt; @@ -9,7 +10,11 @@ use tokio::io::AsyncReadExt; /// There's a better way to implement this. const INITIAL_BUFFER_SIZE : usize = 1024*1024; -pub async fn sora_media_upload_from_file>(file_path: P, creds: CredentialMigrationRef<'_>) -> Result { +pub async fn sora_media_upload_from_file>( + file_path: P, + creds: CredentialMigrationRef<'_>, + maybe_timeout: Option, +) -> Result { let mut file = File::open(&file_path).await?; let mut buffer = Vec::with_capacity(INITIAL_BUFFER_SIZE); file.read_to_end(&mut buffer).await?; @@ -31,7 +36,7 @@ pub async fn sora_media_upload_from_file>(file_path: P, creds: Cr _ => "application/octet-stream", }; - Ok(upload_media_http_request(buffer, filename, mime_type, creds).await?) + Ok(upload_media_http_request(buffer, filename, mime_type, creds, maybe_timeout).await?) } #[cfg(test)] @@ -63,6 +68,7 @@ mod tests { let response = sora_media_upload_from_file( image_path, CredentialMigrationRef::Legacy(&creds), + None, ).await?; println!("media: {:?}", response); diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_http_request.rs b/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_http_request.rs index 722fce788a..3187c28810 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_http_request.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/upload/upload_media_http_request.rs @@ -1,3 +1,4 @@ +use std::time::Duration; use crate::credentials::{SoraCredentials, USER_AGENT}; use crate::creds::credential_migration::CredentialMigrationRef; use crate::sora_error::SoraError; @@ -69,6 +70,7 @@ pub (crate) async fn upload_media_http_request( filename: String, mime_type: &str, credentials: CredentialMigrationRef<'_>, + maybe_timeout: Option, ) -> Result { // Create multipart form @@ -97,11 +99,15 @@ pub (crate) async fn upload_media_http_request( // Make API request let client = Client::new(); - let request_builder = client.post(SORA_UPLOAD_MEDIA_URL) + let mut request_builder = client.post(SORA_UPLOAD_MEDIA_URL) .multipart(form) .header("User-Agent", USER_AGENT) .header("Cookie", &cookie) .header("Authorization", &auth_header); + + if let Some(timeout) = maybe_timeout { + request_builder = request_builder.timeout(timeout); + } let response = request_builder.send().await?;