diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs index 121cfda983..131c5800a4 100644 --- a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs +++ b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs @@ -15,6 +15,7 @@ use openai_sora_client::requests::upload::upload_media_from_file::SoraMediaUploa use std::fs::read_to_string; use std::io::Cursor; use tauri::{AppHandle, Manager, State}; +use openai_sora_client::creds::credential_migration::CredentialMigrationRef; #[tauri::command] pub async fn sora_image_generation_command( @@ -43,7 +44,7 @@ pub async fn generate_image( sora_creds_holder: &SoraCredentialHolder, ) -> AnyhowResult<()> { - let sora_credentials = read_sora_credentials_from_disk(app_data_root) + let sora_credentials_payload = read_sora_credentials_from_disk(app_data_root) .map_err(|err| { error!("Failed to read Sora credentials from disk: {:?}", err); err @@ -58,7 +59,7 @@ pub async fn generate_image( let filename = "image.png".to_string(); - let response = sora_media_upload_from_bytes(image_bytes, filename, &sora_credentials) + let response = sora_media_upload_from_bytes(image_bytes, filename, CredentialMigrationRef::New(&sora_credentials_payload.credentials_set)) .await .map_err(|err| { error!("Failed to upload image to Sora: {:?}", err); @@ -73,7 +74,7 @@ pub async fn generate_image( num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &sora_credentials, + credentials: CredentialMigrationRef::New(&sora_credentials_payload.credentials_set), }).await .map_err(|err| { error!("Failed to call Sora image generation: {:?}", err); diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs index 591e9c478e..2390df9ba2 100644 --- a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs +++ b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs @@ -22,6 +22,7 @@ use std::io::Cursor; use storyteller_client::media_files::get_media_file::get_media_file; use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Manager, State}; +use openai_sora_client::creds::credential_migration::CredentialMigrationRef; use tokens::tokens::media_files::MediaFileToken; #[derive(Deserialize)] @@ -81,14 +82,16 @@ pub async fn generate_image( let files_to_upload = vec![filename]; // TODO(bt,2025-04-21): Read from in-memory cache instead, but allow for desktop replacement. - let mut sora_creds = read_sora_credentials_from_disk(app_data_root) + let sora_creds_payload= read_sora_credentials_from_disk(app_data_root) .map_err(|err| { error!("Failed to read Sora credentials from disk: {:?}", err); err })?; + let mut sora_creds = sora_creds_payload.credentials_set; + let no_sentinel = !app_data_root.get_sora_sentinel_file_path().is_file() - || sora_creds.sentinel.is_none(); + || sora_creds.sora_sentinel.is_none(); if no_sentinel { info!("Refreshing credentials..."); @@ -106,7 +109,7 @@ pub async fn generate_image( let mut sora_media_tokens = Vec::with_capacity(files_to_upload.len()); for file_path in files_to_upload { - let sora_upload_response = sora_media_upload_from_file(file_path, &sora_creds) + let sora_upload_response = sora_media_upload_from_file(file_path, CredentialMigrationRef::New(&sora_creds)) .await?; sora_media_tokens.push(sora_upload_response.id); @@ -121,7 +124,7 @@ pub async fn generate_image( num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &sora_creds, + credentials: CredentialMigrationRef::New(&sora_creds), }).await; if let Err(err) = &response { @@ -141,7 +144,7 @@ pub async fn generate_image( num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &sora_creds, + credentials: CredentialMigrationRef::New(&sora_creds), }).await; } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs b/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs index e356b0ee7c..91591b819a 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/read_sora_credentials_from_disk.rs @@ -1,9 +1,16 @@ use std::fs::read_to_string; use errors::AnyhowResult; use openai_sora_client::credentials::SoraCredentials; +use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; use crate::state::app_dir::AppDataRoot; -pub fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> AnyhowResult { +// TODO: This is only to migrate to the new type. +pub struct CredentialsPayload { + pub credentials: SoraCredentials, + pub credentials_set: SoraCredentialSet, +} + +pub fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> AnyhowResult { let cookie_file = app_data_root.get_sora_cookie_file_path(); let bearer_file = app_data_root.get_sora_bearer_token_file_path(); let sentinel_file = app_data_root.get_sora_sentinel_file_path(); @@ -16,21 +23,29 @@ pub fn read_sora_credentials_from_disk(app_data_root: &AppDataRoot) -> AnyhowRes .trim() .to_string(); + let mut credentials; + if sentinel_file.exists() && sentinel_file.is_file() { let sentinel = read_to_string(sentinel_file)? .trim() .to_string(); - return Ok(SoraCredentials { + credentials = SoraCredentials { bearer_token: bearer, cookie, sentinel: Some(sentinel), - }); + }; + } else { + + credentials = SoraCredentials { + bearer_token: bearer, + cookie, + sentinel: None, + }; } - Ok(SoraCredentials { - bearer_token: bearer, - cookie, - sentinel: None, + Ok(CredentialsPayload { + credentials_set: SoraCredentialSet::from_legacy_credentials(&credentials)?, + credentials, }) } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs index 343e579869..58fc0f859a 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_holder.rs @@ -1,9 +1,8 @@ -use std::ops::Deref; -use std::sync::{Arc, RwLock}; use anyhow::anyhow; use errors::AnyhowResult; use openai_sora_client::credentials::SoraCredentials; use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; +use std::sync::{Arc, RwLock}; // TODO: This is just so we can phase in the new credentials pub struct Inner { @@ -27,11 +26,25 @@ impl SoraCredentialHolder { } } - pub fn set_credentials(&self, credentials: &SoraCredentials) -> AnyhowResult<()> { + pub fn set_legacy_credentials(&self, credentials: &SoraCredentials) -> AnyhowResult<()> { + let credential_set = SoraCredentialSet::from_legacy_credentials(credentials)?; match self.credentials.write() { Err(err) => Err(anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut creds) => { creds.credentials = Some(credentials.clone()); + creds.credential_set = Some(credential_set); + Ok(()) + } + } + } + + pub fn set_credentials(&self, credentials: &SoraCredentialSet) -> AnyhowResult<()> { + let credentials_legacy = credentials.to_legacy_credentials()?; + match self.credentials.write() { + Err(err) => Err(anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut creds) => { + creds.credentials = Some(credentials_legacy); + creds.credential_set = Some(credentials.clone()); Ok(()) } } @@ -42,12 +55,13 @@ impl SoraCredentialHolder { Err(err) => Err(anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut creds) => { creds.credentials = None; + creds.credential_set = None; Ok(()) } } } - pub fn get_credentials(&self) -> AnyhowResult> { + pub fn get_legacy_credentials(&self) -> AnyhowResult> { match self.credentials.read() { Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), Ok(creds) => { @@ -56,7 +70,7 @@ impl SoraCredentialHolder { } } - pub fn get_credentials_required(&self) -> AnyhowResult { + pub fn get_legacy_credentials_required(&self) -> AnyhowResult { match self.credentials.read() { Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), Ok(creds) => match &creds.credentials { @@ -65,4 +79,23 @@ impl SoraCredentialHolder { } } } + + pub fn get_credentials(&self) -> AnyhowResult> { + match self.credentials.read() { + Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(creds) => { + Ok(creds.credential_set.clone()) + } + } + } + + pub fn get_credentials_required(&self) -> AnyhowResult { + match self.credentials.read() { + Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(creds) => match &creds.credential_set { + None => Err(anyhow!("Credentials not set")), + Some(creds) => Ok(creds.clone()), + } + } + } } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs index 4be1ec30c0..a88160daa7 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/sora_credential_manager.rs @@ -9,6 +9,8 @@ use openai_sora_client::requests::sentinel_refresh::generate::token::generate_to use openai_sora_client::requests::sentinel_refresh::refresh_sentinel; use std::fs::OpenOptions; use std::io::Write; +use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; +use openai_sora_client::creds::sora_sentinel::SoraSentinel; #[derive(Clone)] pub struct SoraCredentialManager { @@ -24,7 +26,8 @@ impl SoraCredentialManager { match read_sora_credentials_from_disk(app_data_root) { Err(err) => warn!("Failed to read credentials from disk: {:?}", err), Ok(creds) => { - holder.set_credentials(&creds).expect("Failed to set credentials"); + holder.set_legacy_credentials(&creds.credentials).expect("Failed to set credentials"); + holder.set_credentials(&creds.credentials_set).expect("Failed to set credentials"); } } @@ -34,30 +37,44 @@ impl SoraCredentialManager { } } - pub fn set_credentials(&self, creds: &SoraCredentials) -> AnyhowResult<()> { + pub fn set_legacy_credentials(&self, creds: &SoraCredentials) -> AnyhowResult<()> { + self.holder.set_legacy_credentials(creds) + } + + pub fn set_credentials(&self, creds: &SoraCredentialSet) -> AnyhowResult<()> { self.holder.set_credentials(creds) } pub fn clear_credentials(&self) -> AnyhowResult<()> { - self.holder.clear_credentials() + self.holder.clear_credentials()?; + Ok(()) } - pub fn get_credentials(&self) -> AnyhowResult> { + pub fn get_legacy_credentials(&self) -> AnyhowResult> { + self.holder.get_legacy_credentials() + } + + pub fn get_legacy_credentials_required(&self) -> AnyhowResult { + self.holder.get_legacy_credentials_required() + } + + pub fn get_credentials(&self) -> AnyhowResult> { self.holder.get_credentials() } - pub fn get_credentials_required(&self) -> AnyhowResult { + pub fn get_credentials_required(&self) -> AnyhowResult { self.holder.get_credentials_required() } pub fn reset_from_disk(&self) -> AnyhowResult { let creds = read_sora_credentials_from_disk(&self.app_data_root)?; - self.holder.set_credentials(&creds)?; - Ok(creds) + self.holder.set_credentials(&creds.credentials_set)?; + self.holder.set_legacy_credentials(&creds.credentials)?; + Ok(creds.credentials) } /// Refresh the sentinel token from Sora's API - pub async fn call_sentinel_refresh(&self) -> AnyhowResult { + pub async fn call_sentinel_refresh(&self) -> AnyhowResult { // NB(bt,2025-04-21): Technically we don't need credentials to get a sentinel. let mut creds = self.holder.get_credentials_required()?; @@ -72,7 +89,7 @@ impl SoraCredentialManager { info!("Token obtained."); - creds.sentinel = Some(sentinel.clone()); + creds.sora_sentinel = Some(SoraSentinel::new(sentinel.clone())); self.holder.set_credentials(&creds)?; diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/credential_migration.rs b/crates/lib/api_clients/openai_sora_client/src/creds/credential_migration.rs new file mode 100644 index 0000000000..d578c11571 --- /dev/null +++ b/crates/lib/api_clients/openai_sora_client/src/creds/credential_migration.rs @@ -0,0 +1,12 @@ +use crate::credentials::SoraCredentials; +use crate::creds::sora_credential_set::SoraCredentialSet; + +pub enum CredentialMigration { + Legacy(SoraCredentials), + New(SoraCredentialSet), +} + +pub enum CredentialMigrationRef<'a> { + Legacy(&'a SoraCredentials), + New(&'a SoraCredentialSet), +} diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs b/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs index 685bd191fd..dcd8975f61 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs @@ -1,3 +1,4 @@ +pub mod credential_migration; pub mod sora_cookies; pub mod sora_credential_set; pub mod sora_jwt_bearer_token; diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs index add1f25fde..da4b02fd4a 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs @@ -11,7 +11,11 @@ impl SoraCookies { SoraCookies { cookies } } - pub fn get_cookies(&self) -> &str { + pub fn as_str(&self) -> &str { &self.cookies } + + pub fn to_string(&self) -> String { + self.cookies.clone() + } } diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs index ee00f30da9..2d48e20af4 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs @@ -1,12 +1,15 @@ +use anyhow::anyhow; +use errors::AnyhowResult; +use crate::credentials::SoraCredentials; use crate::creds::sora_cookies::SoraCookies; use crate::creds::sora_jwt_bearer_token::SoraJwtBearerToken; use crate::creds::sora_sentinel::SoraSentinel; #[derive(Clone)] pub struct SoraCredentialSet { - pub (crate) cookies: SoraCookies, - pub (crate) jwt_bearer_token: Option, - pub (crate) sora_sentinel: Option, + pub cookies: SoraCookies, + pub jwt_bearer_token: Option, + pub sora_sentinel: Option, } impl SoraCredentialSet { @@ -23,4 +26,38 @@ impl SoraCredentialSet { SoraCookies::new(cookies.to_string()), ) } + + // TODO(bt,2025-04-23): Just here for migration + pub fn from_legacy_credentials(credentials: &SoraCredentials) -> AnyhowResult { + let cookies = SoraCookies::new(credentials.cookie.clone()); + let jwt_bearer_token = SoraJwtBearerToken::new(credentials.bearer_token.clone())?; + let mut sora_sentinel = None; + + if let Some(sentinel) = &credentials.sentinel { + sora_sentinel = Some(SoraSentinel::new(sentinel.clone())); + } + + Ok(SoraCredentialSet { + cookies, + jwt_bearer_token: Some(jwt_bearer_token), + sora_sentinel, + }) + } + + pub fn to_legacy_credentials(&self) -> AnyhowResult { + let bearer_token= self.jwt_bearer_token + .as_ref() + .map(|token| token.token_str().to_string()) + .ok_or_else(|| anyhow!("There is no bearer token to convert to a required legacy bearer token"))?; + + let sentinel = self.sora_sentinel + .as_ref() + .map(|sentinel| sentinel.get_sentinel().to_string()); + + Ok(SoraCredentials { + cookie: self.cookies.as_str().to_string(), + bearer_token, + sentinel, + }) + } } \ No newline at end of file diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs index 8826f587b8..8176302f82 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs @@ -24,4 +24,11 @@ impl SoraJwtBearerToken { pub fn jwt_claims(&self) -> &SoraJwtClaims { &self.jwt_claims } + + pub fn to_authorization_header_value(&self) -> String { + match self.token.get(0..6) { + Some("bearer") | Some("Bearer") => self.token.clone(), + _ => "Bearer ".to_owned() + &self.token, + } + } } diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs index f73d7a0e5f..40682751ff 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs @@ -3,6 +3,9 @@ #[derive(Clone)] pub struct SoraSentinel { sentinel: String, + + // TODO(bt,2025-04-23): We may want to store the user agent alongside the sentinel + // as that appears to be important for the validity of the sentinel token. } impl SoraSentinel { diff --git a/crates/lib/api_clients/openai_sora_client/src/lib.rs b/crates/lib/api_clients/openai_sora_client/src/lib.rs index 1a9dfe7049..f3c2e26b58 100644 --- a/crates/lib/api_clients/openai_sora_client/src/lib.rs +++ b/crates/lib/api_clients/openai_sora_client/src/lib.rs @@ -1,6 +1,6 @@ pub mod credentials; pub mod creds; +pub mod recipes; pub mod requests; pub mod sora_error; -pub mod utils; -mod recipes; \ No newline at end of file +pub mod utils; \ No newline at end of file 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 aa7bd549b8..cd0b166bf0 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 +1,2 @@ -mod upgrade_or_renew_session; \ No newline at end of file +pub mod 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/upgrade_or_renew_session.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/upgrade_or_renew_session.rs index 38e1cc1af3..7af04831d5 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/upgrade_or_renew_session.rs +++ b/crates/lib/api_clients/openai_sora_client/src/recipes/upgrade_or_renew_session.rs @@ -30,7 +30,7 @@ pub async fn upgrade_or_renew_session(sora_credentials: &mut SoraCredentialSet) } if refresh_jwt { - let cookies = sora_credentials.cookies.get_cookies(); + let cookies = sora_credentials.cookies.as_str(); info!("Generating new JWT bearer token..."); let token = generate_bearer_with_cookie(cookies).await?; 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 new file mode 100644 index 0000000000..dbf88733dc --- /dev/null +++ b/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs @@ -0,0 +1,51 @@ +use crate::credentials::SoraCredentials; +use crate::requests::image_gen::image_gen_status::TaskStatus; +use crate::requests::image_gen::{get_image_gen_status, StatusRequest, TaskResponse, VideoGenStatusResponse}; +use errors::AnyhowResult; + +pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCredentials, retry_limit: Option) -> AnyhowResult { + let status_request = StatusRequest { + limit: None, + before: None, + // before: Some(task_id.clone()), + }; + let retry_limit = retry_limit.unwrap_or(10); + + 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?; + + match find_task_response_by_id(&status_response, task_id.clone()) { + Some(task_response) => match TaskStatus::from_str(&task_response.status) { + TaskStatus::Succeeded => { + return Ok(task_response.clone()); + }, + TaskStatus::Failed => { + return Err(anyhow::anyhow!("Task failed")); + }, + TaskStatus::Unknown(status) => { + println!("Unknown task status: {}", status); + continue; + }, + TaskStatus::Queued => { + println!("Task is queued"); + continue; + }, + TaskStatus::Running => { + println!("Task is running"); + continue; + }, + }, + None => { + return Err(anyhow::anyhow!("Task not found")); + }, + } + } + + Err(anyhow::anyhow!("Task not found")) +} + +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) +} 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 d0a1d39750..67992c746c 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,5 +1,5 @@ +use crate::creds::credential_migration::CredentialMigrationRef; use serde_derive::{Deserialize, Serialize}; -use crate::credentials::SoraCredentials; use thiserror::Error; const SORA_IMAGE_GEN_URL: &str = "https://sora.com/backend/video_gen"; @@ -146,20 +146,41 @@ 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: &SoraCredentials) -> Result { +pub (crate) async fn image_gen_http_request(sora_request: RawSoraImageGenRequest, credentials: CredentialMigrationRef<'_>) -> Result { let client = reqwest::Client::new(); - let sentinel = match credentials.sentinel.as_deref() { - None => return Err(SoraError::GenericError("Calls to image gen require a sentinel in the credentials payload.".to_string())), - Some(sentinel) => sentinel, - }; + let mut cookie; + let mut authorization_header; + let mut sentinel; + + match credentials { + CredentialMigrationRef::Legacy(creds) => { + cookie = creds.cookie.clone(); + authorization_header = creds.authorization_header_value(); + // TODO(bt,2025-04-23): We're using a Sora payload error in place of application state error. Surface this differently. + sentinel = creds.sentinel.as_ref() + .map(|sentinel| sentinel.to_string()) + .ok_or(SoraError::SentinelBlock("Sentinel is required for image generation.".to_string()))?; + } + CredentialMigrationRef::New(creds) => { + cookie = creds.cookies.to_string(); + // TODO(bt,2025-04-23): We're using a Sora payload error in place of application state error. Surface this differently. + authorization_header = creds.jwt_bearer_token.as_ref() + .ok_or(SoraError::InvalidJwt("JWT bearer is required for image generation".to_string()))? + .to_authorization_header_value(); + // TODO(bt,2025-04-23): We're using a Sora payload error in place of application state error. Surface this differently. + sentinel = creds.sora_sentinel.as_ref() + .map(|sentinel| sentinel.get_sentinel().to_string()) + .ok_or(SoraError::SentinelBlock("Sentinel is required for image generation.".to_string()))?; + } + } let http_request = client.post(SORA_IMAGE_GEN_URL) .header("User-Agent", USER_AGENT) - .header("Cookie", &credentials.cookie) - .header("Authorization", credentials.authorization_header_value()) + .header("Cookie", &cookie) + .header("Authorization", &authorization_header) .header("Content-Type", "application/json") - .header("OpenAI-Sentinel-Token", sentinel); + .header("OpenAI-Sentinel-Token", &sentinel); let http_request = http_request.json(&sora_request).build() .map_err(|e| SoraError::NetworkError(e.to_string()))?; 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 d4ce76092b..f1012da827 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 @@ -220,47 +220,6 @@ pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: & Ok(response) } -pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCredentials, retry_limit: Option) -> AnyhowResult { - let status_request = StatusRequest { - limit: None, - before: None, - // before: Some(task_id.clone()), - }; - let retry_limit = retry_limit.unwrap_or(10); - - 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?; - match find_task_response_by_id(&status_response, task_id.clone()) { - Some(task_response) => match TaskStatus::from_str(&task_response.status) { - TaskStatus::Succeeded => { - return Ok(task_response.clone()); - }, - TaskStatus::Failed => { - return Err(anyhow::anyhow!("Task failed")); - }, - TaskStatus::Unknown(status) => { - println!("Unknown task status: {}", status); - continue; - }, - TaskStatus::Queued => { - println!("Task is queued"); - continue; - }, - TaskStatus::Running => { - println!("Task is running"); - continue; - }, - }, - None => { - return Err(anyhow::anyhow!("Task not found")); - }, - } - } - - Err(anyhow::anyhow!("Task not found")) -} - pub async fn save_generations_to_dir(generations: &[Generation], dir: &str) -> AnyhowResult<()> { let dir = Path::new(dir); if !dir.exists() { @@ -288,14 +247,11 @@ pub async fn save_generations_to_dir(generations: &[Generation], dir: &str) -> A Ok(()) } -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) -} - #[cfg(test)] mod tests { use crate::credentials::SoraCredentials; - use crate::image_gen::image_gen_status::{get_image_gen_status, save_generations_to_dir, wait_for_image_gen_status, StatusRequest, VideoGenStatusResponse}; + use crate::recipes::wait_for_image_gen_status::wait_for_image_gen_status; + use crate::requests::image_gen::image_gen_status::{get_image_gen_status, save_generations_to_dir, StatusRequest, VideoGenStatusResponse}; use errors::AnyhowResult; use std::fs::read_to_string; use testing::test_file_path::test_file_path; diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs index 525b0fbf62..fb9c54de2f 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs @@ -4,4 +4,4 @@ pub mod common; pub mod sora_image_gen_remix; pub mod sora_image_gen_simple; pub mod image_gen_status; -pub use image_gen_status::{get_image_gen_status, save_generations_to_dir, wait_for_image_gen_status, Generation, StatusRequest, TaskResponse, VideoGenStatusResponse}; +pub use image_gen_status::{get_image_gen_status, save_generations_to_dir, Generation, StatusRequest, TaskResponse, VideoGenStatusResponse}; 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 b5520b012d..c0c4f9c13d 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,4 +1,4 @@ -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, InpaintItem, InpaintItemType, OperationType, RawSoraImageGenRequest, SoraError, VideoGenType}; @@ -7,7 +7,7 @@ pub struct SoraImageGenRemixRequest<'a> { pub num_images: NumImages, pub image_size: ImageSize, pub sora_media_tokens: Vec, - pub credentials: &'a SoraCredentials, + pub credentials: CredentialMigrationRef<'a>, } /// The "remix" commands let you supply additional images as context. @@ -35,7 +35,7 @@ pub async fn sora_image_gen_remix(request: SoraImageGenRemixRequest<'_>) -> Resu }).collect(), }; - let result = image_gen_http_request(args, &request.credentials).await?; + let result = image_gen_http_request(args, request.credentials).await?; Ok(SoraImageGenResponse { task_id: result.id, @@ -45,6 +45,7 @@ pub async fn sora_image_gen_remix(request: SoraImageGenRemixRequest<'_>) -> Resu #[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_remix::{sora_image_gen_remix, SoraImageGenRemixRequest}; use errors::AnyhowResult; @@ -74,7 +75,7 @@ mod tests { num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: vec!["media_01jqyhrz4detyvtzwp2p4j63ad".to_string()], - credentials: &creds, + credentials: CredentialMigrationRef::Legacy(&creds), }).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 4cc710aa93..38c0f60bab 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 @@ -2,12 +2,13 @@ use crate::credentials::SoraCredentials; 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; pub struct SoraImageGenSimpleRequest<'a> { pub prompt: String, pub num_images: NumImages, pub image_size: ImageSize, - pub credentials: &'a SoraCredentials, + pub credentials: CredentialMigrationRef<'a>, } pub async fn sora_image_gen_simple(args: SoraImageGenSimpleRequest<'_>) -> AnyhowResult { @@ -61,7 +62,7 @@ mod tests { prompt: "A pirate and a ninja fight in a battle inside a UFO. Fully photo realistic, lifelike, lens flare".to_string(), num_images: NumImages::One, image_size: ImageSize::Square, - credentials: &creds, + credentials: creds, }).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 1f450f1778..ba0d8b0d9b 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 @@ -1,16 +1,11 @@ -use crate::credentials::SoraCredentials; +use crate::creds::credential_migration::CredentialMigrationRef; use crate::requests::upload::upload_media_http_request::{upload_media_http_request, SoraMediaUploadResponse}; use errors::AnyhowResult; use std::path::PathBuf; -pub struct SoraMediaUploadRequest<'a> { - pub file_path: String, - pub credentials: &'a SoraCredentials, -} - /// 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: &SoraCredentials) -> AnyhowResult { +pub async fn sora_media_upload_from_bytes(bytes: Vec, file_name: String, creds: CredentialMigrationRef<'_>) -> AnyhowResult { let file_path = PathBuf::from(&file_name); // TODO: Read file magic bytes first, then fall back to this. 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 3b1ed6c4fa..df6c8981aa 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 @@ -4,6 +4,7 @@ use errors::AnyhowResult; use std::path::Path; use tokio::fs::File; use tokio::io::AsyncReadExt; +use crate::creds::credential_migration::CredentialMigrationRef; /// Try to prevent buffer reallocations. /// There's a better way to implement this. @@ -14,7 +15,7 @@ pub struct SoraMediaUploadRequest<'a> { pub credentials: &'a SoraCredentials, } -pub async fn sora_media_upload_from_file>(file_path: P, creds: &SoraCredentials) -> AnyhowResult { +pub async fn sora_media_upload_from_file>(file_path: P, creds: CredentialMigrationRef<'_>) -> AnyhowResult { let mut file = File::open(&file_path).await?; let mut buffer = Vec::with_capacity(INITIAL_BUFFER_SIZE); file.read_to_end(&mut buffer).await?; @@ -42,10 +43,11 @@ pub async fn sora_media_upload_from_file>(file_path: P, creds: &S #[cfg(test)] mod tests { use crate::credentials::SoraCredentials; - use crate::upload::upload_media_from_file::sora_media_upload_from_file; + use crate::requests::upload::upload_media_from_file::sora_media_upload_from_file; 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] @@ -66,7 +68,7 @@ mod tests { let response = sora_media_upload_from_file( image_path, - &creds, + CredentialMigrationRef::Legacy(&creds), ).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 be69f14e26..ff567ad099 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,4 +1,5 @@ use crate::credentials::{SoraCredentials, USER_AGENT}; +use crate::creds::credential_migration::CredentialMigrationRef; use crate::sora_error::SoraError; use reqwest::multipart::{Form, Part}; use reqwest::Client; @@ -65,7 +66,7 @@ pub (crate) async fn upload_media_http_request( file_bytes: Vec, filename: String, mime_type: &str, - credentials: &SoraCredentials, + credentials: CredentialMigrationRef<'_>, ) -> Result { // Create multipart form @@ -75,13 +76,30 @@ pub (crate) async fn upload_media_http_request( let form = Form::new().part("file", part); + let cookie; + let auth_header; + + match credentials { + CredentialMigrationRef::Legacy(creds) => { + cookie = creds.cookie.clone(); + auth_header = creds.authorization_header_value(); + } + CredentialMigrationRef::New(creds) => { + cookie = creds.cookies.to_string(); + auth_header = creds.jwt_bearer_token + .as_ref() + .ok_or(SoraError::NoBearerTokenAvailable)? + .to_authorization_header_value(); + } + } + // Make API request let client = Client::new(); let request_builder = client.post(SORA_UPLOAD_MEDIA_URL) .multipart(form) .header("User-Agent", USER_AGENT) - .header("Cookie", &credentials.cookie) - .header("Authorization", credentials.authorization_header_value()); + .header("Cookie", &cookie) + .header("Authorization", &auth_header); let response = request_builder.send().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 7266b86595..aadd03ed74 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 @@ -3,6 +3,9 @@ use std::fmt::{Display, Formatter}; #[derive(Debug)] pub enum SoraError { + /// We haven't received a bearer token yet. + NoBearerTokenAvailable, + /// Unauthorized, cookie and/or bearer token expired. We'll need to ask for a refreshed login. /// /// Example message, e.g. from the upload endpoint: @@ -28,13 +31,16 @@ impl Error for SoraError {} impl Display for SoraError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { - SoraError::UnauthorizedCookieOrBearerExpired => { + Self::NoBearerTokenAvailable => { + write!(f, "Unauthorized: no bearer token set") + } + Self::UnauthorizedCookieOrBearerExpired => { write!(f, "Unauthorized: cookie and/or bearer token expired") } - SoraError::OtherBadStatus(err) => { + Self::OtherBadStatus(err) => { write!(f, "Other error: {}", err) } - SoraError::ReqwestError(err) => { + Self::ReqwestError(err) => { write!(f, "Reqwest error: {}", err) } } diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs index d55cd4d1ce..9e348e4c13 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs @@ -28,8 +28,9 @@ use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes; use mysql_queries::queries::batch_generations::insert_batch_generation_records::{insert_batch_generation_records, InsertBatchArgs}; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use openai_sora_client::credentials::SoraCredentials; +use openai_sora_client::recipes::wait_for_image_gen_status::wait_for_image_gen_status; use openai_sora_client::requests::image_gen::image_gen_status::TaskStatus; -use openai_sora_client::requests::image_gen::{wait_for_image_gen_status, Generation}; +use openai_sora_client::requests::image_gen::Generation; use r2d2_redis::r2d2::PooledConnection; use r2d2_redis::redis::Commands; use r2d2_redis::RedisConnectionManager; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs index 1c908754ca..11b02197aa 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs @@ -37,6 +37,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::{Ge use mysql_queries::payloads::generic_inference_args::inner_payloads::sora_image_gen_args::SoraImageGenArgs; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use openai_sora_client::creds::credential_migration::CredentialMigrationRef; 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_file::{sora_media_upload_from_file, SoraMediaUploadRequest}; @@ -268,7 +269,7 @@ pub async fn enqueue_studio_image_generation_handler(http_request: HttpRequest, let mut sora_media_tokens = Vec::with_capacity(files_to_upload.len()); for file_path in files_to_upload { - let sora_upload_response = sora_media_upload_from_file(file_path, &sora_credentials).await.map_err(|err| { + let sora_upload_response = sora_media_upload_from_file(file_path, CredentialMigrationRef::Legacy(&sora_credentials)).await.map_err(|err| { error!("Failed to upload scene media to Sora: {:?}", err); EnqueueImageGenRequestError::ServerError })?; @@ -299,7 +300,7 @@ pub async fn enqueue_studio_image_generation_handler(http_request: HttpRequest, num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &sora_credentials + credentials: CredentialMigrationRef::Legacy(&sora_credentials), }).await; debug!("Sora image gen response: {:?}", response); @@ -328,7 +329,7 @@ pub async fn enqueue_studio_image_generation_handler(http_request: HttpRequest, num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &updated_sora_credentials + credentials: CredentialMigrationRef::Legacy(&updated_sora_credentials) }).await; }, Err(e) => { @@ -347,7 +348,7 @@ pub async fn enqueue_studio_image_generation_handler(http_request: HttpRequest, num_images: NumImages::One, image_size: ImageSize::Square, sora_media_tokens: sora_media_tokens.clone(), - credentials: &updated_sora_credentials + credentials: CredentialMigrationRef::Legacy(&updated_sora_credentials) }).await; }, Err(e) => {