improve sora status checking errors

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)))
 }

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<u32>) -> AnyhowResult<TaskResponse> {
   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) {

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<VideoGenStatusResponse, SoraError> {
+pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentialSet) -> Result<VideoGenStatusResponse, SoraError> {
+  
+  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::<VideoGenStatusResponse>(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();

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)
+      }
     }
   }
 }