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 index b25a8219c3..b7abbc5655 100644 --- 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 @@ -39,11 +39,10 @@ pub async fn list_sora_task_status_with_session_auto_renew( new_creds = Some(updated_creds); } - let legacy_creds = new_creds.as_ref() - .unwrap_or_else(|| &args.credentials) - .to_legacy_credentials()?; + let upgraded_creds = new_creds.as_ref() + .unwrap_or_else(|| &args.credentials); - let result = get_image_gen_status(&request, &legacy_creds).await; + let result = get_image_gen_status(&request, &upgraded_creds).await; let err = match result { Ok(response) => return Ok((response, new_creds)), @@ -56,7 +55,7 @@ pub async fn list_sora_task_status_with_session_auto_renew( info!("Retrying image upload with new credentials..."); - let result = get_image_gen_status(&request, &legacy_creds).await?; + let result = get_image_gen_status(&request, &upgraded_creds).await?; Ok((result, Some(new_creds))) } 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 50d5bf0b74..faa2051dc0 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 @@ -2,6 +2,7 @@ use crate::credentials::SoraCredentials; use crate::requests::image_gen::image_gen_status::TaskStatus; use crate::requests::image_gen::image_gen_status::{get_image_gen_status, StatusRequest, TaskResponse, VideoGenStatusResponse}; use errors::AnyhowResult; +use crate::creds::sora_credential_set::SoraCredentialSet; pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCredentials, retry_limit: Option) -> AnyhowResult { let status_request = StatusRequest { @@ -10,11 +11,13 @@ pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCrede // before: Some(task_id.clone()), }; let retry_limit = retry_limit.unwrap_or(10); + + let new_creds = SoraCredentialSet::from_legacy_credentials(&credentials)?; for _ in 0..retry_limit { tokio::time::sleep(std::time::Duration::from_secs(10)).await; - let status_response = get_image_gen_status(&status_request, credentials).await?; + let status_response = get_image_gen_status(&status_request, &new_creds).await?; match find_task_response_by_id(&status_response, task_id.clone()) { Some(task_response) => match TaskStatus::from_str(&task_response.status) { 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 1411f69670..dd47760fec 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 @@ -6,7 +6,9 @@ use reqwest::Url; use serde_derive::{Deserialize, Serialize}; use std::fmt::{Debug, Display}; use std::{fs::File, io::Write, path::Path}; -use log::debug; +use log::{debug, warn}; +use crate::creds::sora_credential_set::SoraCredentialSet; +use crate::creds::sora_jwt_bearer_token::SoraJwtBearerToken; const SORA_STATUS_URL: &str = "https://sora.com/backend/video_gen"; @@ -199,7 +201,16 @@ impl StatusRequest { } /// Gets the status of image generation tasks from Sora API -pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentials) -> Result { +pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentialSet) -> Result { + + let bearer_header = match credentials.jwt_bearer_token.as_ref() { + Some(bearer) => bearer.to_authorization_header_value(), + None => { + warn!("No JWT bearer for getting image gen status!"); + return Err(SoraError::NoBearerTokenAvailable) + }, + }; + let client = reqwest::Client::new(); let mut url = reqwest::Url::parse(SORA_STATUS_URL) @@ -216,8 +227,8 @@ pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: & let http_request = client.get(url) .header("User-Agent", USER_AGENT) - .header("Cookie", &credentials.cookie) - .header("Authorization", credentials.authorization_header_value()) + .header("Cookie", credentials.cookies.as_str()) + .header("Authorization", bearer_header) .header("Content-Type", "application/json"); let http_request = http_request.build()?; @@ -235,9 +246,15 @@ pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: & return Err(SoraError::AnyhowError(anyhow::anyhow!("The status request failed; status = {:?}, message = {:?}", status, response_body))); } - let response: VideoGenStatusResponse = serde_json::from_str(response_body)?; + let maybe_response = serde_json::from_str::(response_body); - Ok(response) + match maybe_response { + Ok(response) => Ok(response), + Err(err) => { + warn!("Failed to parse status response: {:?}", err); + Err(SoraError::JsonErrorWithBody(err, response_body.to_string())) + } + } } pub async fn save_generations_to_dir(generations: &[Generation], dir: &str) -> AnyhowResult<()> { @@ -275,6 +292,7 @@ mod tests { use errors::AnyhowResult; use std::fs::read_to_string; use testing::test_file_path::test_file_path; + use crate::creds::sora_credential_set::SoraCredentialSet; #[ignore] #[tokio::test] @@ -297,10 +315,12 @@ mod tests { let bearer = bearer.trim().to_string(); let creds = SoraCredentials { bearer_token: bearer, cookie, sentinel: Some(sentinel) }; + let new_creds = SoraCredentialSet::from_legacy_credentials(&creds)?; + // Get the task status for a specific task // let response = get_image_gen_status(&StatusRequest::new(None, Some("task_01jr9yvpfyetx9r7qvvx38scna".to_string())), &creds).await?; - let response = get_image_gen_status(&StatusRequest::new(Some(50), None), &creds).await?; + let response = get_image_gen_status(&StatusRequest::new(Some(50), None), &new_creds).await?; println!("task_id: {}", response.task_responses[0].id); let task_id = response.task_responses[0].id.clone(); 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 d8edda998c..897df7ff75 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 @@ -62,6 +62,10 @@ pub enum SoraError { /// serde_json::Error, likely from JSON deserialization schema mismatch. JsonError(serde_json::Error), + + /// serde_json::Error, likely from JSON deserialization schema mismatch. + /// Includes the original body. + JsonErrorWithBody(serde_json::Error, String), } impl Error for SoraError {} @@ -108,6 +112,9 @@ impl Display for SoraError { Self::JsonError(err) => { write!(f, "serde_json error: {}", err) } + Self::JsonErrorWithBody(err, message) => { + write!(f, "serde_json error: {:?} body: {:?}", err, message) + } } } }