diff --git a/crates/api_clients/grok_client/src/error/grok_client_error.rs b/crates/api_clients/grok_client/src/error/grok_client_error.rs index 52f419f1e..6c1b52157 100644 --- a/crates/api_clients/grok_client/src/error/grok_client_error.rs +++ b/crates/api_clients/grok_client/src/error/grok_client_error.rs @@ -56,6 +56,9 @@ pub enum GrokClientError { /// Can't make request because cookies aren't present NoCookiesPresent, + /// Unknown error generating video + ErrorGeneratingVideo, + // /// An error reading the file for upload. // FileForUploadReadError(std::io::Error), // @@ -115,6 +118,7 @@ impl Display for GrokClientError { Self::BadSignatureInputs => write!(f, "Bad signature inputs"), Self::TimeoutMathBroken => write!(f, "Timeout math is broken"), Self::NoCookiesPresent => write!(f, "No cookies present"), + Self::ErrorGeneratingVideo => write!(f, "Error generating video"), //Self::FileForUploadReadError(err) => write!(f, "Error reading file for upload: {}", err), //Self::FileForUploadHasInvalidPath => write!(f, "The file path provided for upload is invalid."), diff --git a/crates/api_clients/grok_client/src/recipes/mod.rs b/crates/api_clients/grok_client/src/recipes/mod.rs index e774ecd05..68bb01157 100644 --- a/crates/api_clients/grok_client/src/recipes/mod.rs +++ b/crates/api_clients/grok_client/src/recipes/mod.rs @@ -1,3 +1,4 @@ +pub mod prompt_websocket_image_with_retry; pub mod request_client_secrets; pub mod upload_image_and_generate_video; pub mod upload_image_and_generate_video_with_retry; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/recipes/prompt_websocket_image_with_retry.rs b/crates/api_clients/grok_client/src/recipes/prompt_websocket_image_with_retry.rs new file mode 100644 index 000000000..8d2d90ee0 --- /dev/null +++ b/crates/api_clients/grok_client/src/recipes/prompt_websocket_image_with_retry.rs @@ -0,0 +1,189 @@ +use crate::credentials::grok_cookies::GrokCookies; +use crate::error::grok_client_error::GrokClientError; +use crate::error::grok_error::GrokError; +use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; +use crate::requests::image_websocket::grok_websocket::GrokWebsocket; +use crate::requests::image_websocket::messages::websocket_client_message::{ClientMessageAspectRatio, WebsocketClientMessage}; +use log::{error, info}; + +pub struct PromptWebsocketImageWithRetryArgs<'a> { + pub websocket: &'a mut GrokWebsocket, + pub prompt: &'a str, + pub aspect_ratio: ClientMessageAspectRatio, + pub cookies: &'a GrokCookies, +} + +pub struct PromptWebsocketImageWithRetryResult { + pub maybe_new_websocket: Option, +} + + +pub async fn prompt_websocket_image_with_retry(args: PromptWebsocketImageWithRetryArgs<'_>) -> Result +{ + let mut maybe_new_websocket : Option = None; + + for i in 0..3 { + info!("Starting Grok Websocket image generation (attempt {}) ...", (i+1)); + + let message = WebsocketClientMessage::new_image_prompt(args.prompt, args.aspect_ratio); + + let result = + if let Some(websocket) = maybe_new_websocket.as_mut() { + websocket.send_serializable(&message).await + } else { + args.websocket.send_serializable(&message).await + }; + + match result { + Ok(()) => { + break; + } + Err(GrokError::Client(GrokClientError::WreqClientError(err))) => { + error!("Error sending message to websocket (WreqClientError): {}", err); + } + Err(GrokError::Client(err)) => { + error!("Error sending message to websocket (GrokClientError): {}", err); + } + Err(GrokError::ApiSpecific(err)) => { + error!("Error sending message to websocket (GrokSpecificApiError): {}", err); + } + Err(GrokError::ApiGeneric(err)) => { + error!("Error sending message to websocket (GrokGenericApiError): {}", err); + } + } + + let result = create_listen_websocket(CreateListenWebsocketArgs { + cookies: args.cookies.as_str(), + }).await?; + + maybe_new_websocket = Some(GrokWebsocket::new(result)); + } + + Ok(PromptWebsocketImageWithRetryResult { + maybe_new_websocket, + }) +} + +#[cfg(test)] +mod tests { + use crate::recipes::prompt_websocket_image_with_retry::{prompt_websocket_image_with_retry, PromptWebsocketImageWithRetryArgs}; + use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; + use crate::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; + use crate::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; + use crate::requests::image_websocket::messages::websocket_server_message::WebsocketServerMessage; + use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; + use crate::test_utils::get_test_cookies::{get_test_cookies, get_typed_test_cookies}; + use errors::AnyhowResult; + use futures::SinkExt; + use log::warn; + use std::io::Write; + use std::time::Duration; + use wreq::ws::message::{CloseCode, Message, Utf8Bytes}; + + #[tokio::test] + #[ignore] // manually test + async fn prompt() -> AnyhowResult<()> { + //setup_test_logging(LevelFilter::Trace); + + let cookies = get_typed_test_cookies()?; + + let websocket = create_listen_websocket(CreateListenWebsocketArgs { + cookies: cookies.as_str(), + }).await?; + + //let websocket = GrokWrappedWebsocket::new(websocket); + let mut websocket = GrokWebsocket::new(websocket); + + //println!("Sending..."); + //std::io::stdout().flush()?; + + //println!("For our test, we'll close the websocket and retry..."); + //websocket.websocket + // .close(CloseCode::NORMAL, Utf8Bytes::try_from(Vec::new())?).await?; + + //let mut websocket = GrokWebsocket::new(websocket.websocket); + + let result = prompt_websocket_image_with_retry(PromptWebsocketImageWithRetryArgs { + websocket: &mut websocket, + prompt: "a dog riding a motorcycle", + aspect_ratio: ClientMessageAspectRatio::WideThreeByTwo, + cookies: &cookies, + }).await?; + + println!("Reading..."); + std::io::stdout().flush()?; + + let mut count = 0; + + println!("Polling."); + + loop { + //let maybe_message = + // websocket.get_response_with_timeout(Duration::from_millis(1000)).await?; + + let maybe_message = + websocket.try_next_timeout(Duration::from_millis(1000)).await?; + + let mut maybe_raw_text : Option = None; + + let maybe_message = match maybe_message { + None => None, + Some(Message::Text(text)) => { + let maybe_message = WebsocketServerMessage::from_json_str(text.as_str())?; + let mut text = text.as_str().to_string(); + text.truncate(100); + //println!("Received message: {}", text); + maybe_raw_text = Some(text); + Some(maybe_message) + }, + Some(_) => { + warn!("Received non-text websocket message."); + None + }, + }; + + match maybe_message { + None => { + println!("No message received within timeout."); + count = count + 1; + } + Some(message) => { + match message { + WebsocketServerMessage::Image(image) => { + println!("IMAGE: {:?}", image.percentage_complete); + } + WebsocketServerMessage::Json(json) => { + println!("JSON : {:?}", json.percentage_complete); + } + WebsocketServerMessage::Unknown(unknown) => { + //let typ = unknown.get("type"); + let unknown_string = unknown.as_str(); + println!("[UNKNOWN] websocket message: {:?}", maybe_raw_text); + } + } + count = 0; + } + } + + // TODO: Make sure we capture all the events for images + // TODO: Explore the other APIs + // TODO understand the IDs + // TODO: Video websocket APIs + // TODO: Upload APIs + + if count > 5 { + println!("No messages after 5 seconds"); + break; + } + } + + println!("Done polling."); + + log::logger().flush(); + + assert_eq!(1,2); + + Ok(()) + } +} diff --git a/crates/api_clients/grok_client/src/recipes/upload_image_and_generate_video_with_retry.rs b/crates/api_clients/grok_client/src/recipes/upload_image_and_generate_video_with_retry.rs index b7f336515..4e981563d 100644 --- a/crates/api_clients/grok_client/src/recipes/upload_image_and_generate_video_with_retry.rs +++ b/crates/api_clients/grok_client/src/recipes/upload_image_and_generate_video_with_retry.rs @@ -1,7 +1,22 @@ +use crate::credentials::grok_client_secrets::GrokClientSecrets; +use crate::credentials::grok_full_credentials::GrokFullCredentials; +use crate::datatypes::api::aspect_ratio::AspectRatio; +use crate::datatypes::file_upload_spec::FileUploadSpec; +use crate::error::grok_client_error::GrokClientError; +use crate::error::grok_error::GrokError; +use crate::error::grok_generic_api_error::GrokGenericApiError; +use crate::error::grok_specific_api_error::GrokSpecificApiError; +use crate::recipes::request_client_secrets::{request_client_secrets, RequestClientSecretsArgs}; +use crate::recipes::upload_image_and_generate_video::{ImageUploadAndGenerateVideoResult, UploadImageAndGenerateVideo}; +use crate::requests::media_posts::create_media_post::grok_create_media_post::{GrokCreateMediaPost, MediaPostType}; +use crate::requests::media_posts::like_media_post::grok_like_media::GrokLikeMediaPost; +use crate::requests::upload_file::grok_upload_file::GrokUploadFile; +use crate::requests::video_chat::grok_video_gen_chat_conversation::{GrokVideoGenChatConversationBuilder, VideoMediaPostType}; +use crate::utils::user_and_file_id_to_image_url::user_and_file_id_to_image_url; +use crate::utils::user_and_file_id_to_video_url::user_and_file_id_to_video_url; +use log::{error, info}; use std::path::Path; use std::time::Duration; -use crate::credentials::grok_full_credentials::GrokFullCredentials; -use crate::datatypes::file_upload_spec::FileUploadSpec; pub struct UploadImageAndGenerateVideoWithRetry<'a, P: AsRef> { pub credentials: &'a GrokFullCredentials, @@ -12,6 +27,9 @@ pub struct UploadImageAndGenerateVideoWithRetry<'a, P: AsRef> { /// Video generation prompt pub prompt: Option<&'a str>, + /// Aspect ratio for the video. + pub aspect_ratio: Option, + /// Wait for the full video to be generated before returning /// By default, the endpoint stays open for 30-ish seconds while /// the video generates and streams back JSON progress updates. @@ -22,3 +40,302 @@ pub struct UploadImageAndGenerateVideoWithRetry<'a, P: AsRef> { pub individual_request_timeout: Option, } + +#[derive(Clone)] +pub struct ImageUploadAndGenerateVideoWithRetryResult { + /// If new secrets were generated, they're stored here. + pub maybe_new_client_secrets: Option, + + /// The video generation result. + pub upload_result: ImageUploadAndGenerateVideoResult, +} + +pub async fn upload_image_and_generate_video_with_retry>( + args: UploadImageAndGenerateVideoWithRetry<'_, P> +) -> Result { + + let mut current_full_credentials_ref = args.credentials; + let mut maybe_new_credentials = None; + + info!("Uploading file to Grok..."); + + let request = GrokUploadFile { + file: args.file, + cookie: current_full_credentials_ref.cookies.to_string(), + request_timeout: args.individual_request_timeout, + }; + + let upload_result = request.upload().await?; + + info!("File URI: {:?}", upload_result.file_uri); + info!("File Metadata ID: {:?}", upload_result.file_id); + + let upload_file_id = match upload_result.file_id { + Some(id) => id, + None => { + error!("Previous file upload failed. Cannot continue."); + return Err(GrokGenericApiError::UploadFailed.into()); + } + }; + + let url = user_and_file_id_to_image_url( + ¤t_full_credentials_ref.client_secrets.user_id, + &upload_file_id + ); + + info!("Uploaded URI: {:?}", url); + + info!("Creating media post..."); + + let request = GrokCreateMediaPost { + user_id: ¤t_full_credentials_ref.client_secrets.user_id, + file_id: &upload_file_id, + media_type: MediaPostType::UserUploadedImage, + cookie: current_full_credentials_ref.cookies.as_str(), + request_timeout: args.individual_request_timeout, + }; + + let post_result = request.send().await?; + + info!("Posted Image."); + + info!("Video prompt: {:?}", args.prompt); + info!("Video aspect ratio: {:?}", args.aspect_ratio); + + let mut video_enqueued_successfully = false; + let mut generation_is_complete= false; + let mut maybe_video_file_id = None; + let mut last_error = None; + + for i in 0..3 { + info!("Generate video... attempt {}", (i+1)); + + let request = GrokVideoGenChatConversationBuilder { + user_id: ¤t_full_credentials_ref.client_secrets.user_id, + file_id: &upload_file_id, + media_type: VideoMediaPostType::UserUploadedImage, + cookie: current_full_credentials_ref.cookies.as_str(), + prompt: args.prompt, + aspect_ratio: args.aspect_ratio, + request_timeout: args.individual_request_timeout, + baggage: ¤t_full_credentials_ref.client_secrets.baggage, + sentry_trace: ¤t_full_credentials_ref.client_secrets.sentry_trace, + verification_token: ¤t_full_credentials_ref.client_secrets.verification_token, + svg_data: ¤t_full_credentials_ref.client_secrets.svg_path_data, + numbers: ¤t_full_credentials_ref.client_secrets.numbers, + }; + + let video_gen_result = + if args.wait_for_generation { + request.wait_for_video() + .await + .map(|res| res.video_file_id) + } else { + request.stream_only_video_id() + .await + .map(|res| res.video_file_id) + }; + + match video_gen_result { + Ok(res) => { + info!("Video Generation Enqueued!"); + video_enqueued_successfully = true; + generation_is_complete = args.wait_for_generation; // If we synchronously waited, it'll be complete. + maybe_video_file_id = res; + break; + } + Err(GrokError::ApiSpecific(GrokSpecificApiError::AutomationBlocked)) => { + info!("Grok automation blocked; renewing credentials."); + last_error = Some(GrokError::ApiSpecific(GrokSpecificApiError::AutomationBlocked)); + } + Err(err) => { + info!("Grok error; renewing credentials. Error = {:?}", err); + last_error = Some(err); + } + } + + info!("Refreshing Grok client secrets..."); + + let secrets = request_client_secrets(RequestClientSecretsArgs { + cookies: ¤t_full_credentials_ref.cookies, + }).await?; + + // NB: This is just to appease the borrow checker. It doesn't like that we're borrowing `maybe_new_credentials` + // in a loop and replacing it while borrowed, hence this hack. + current_full_credentials_ref = args.credentials; + + maybe_new_credentials = Some(GrokFullCredentials::from_cookies_and_client_secrets( + current_full_credentials_ref.cookies.clone(), + secrets, + )); + + if let Some(creds) = maybe_new_credentials.as_ref() { + current_full_credentials_ref = creds; + } + } + + if !video_enqueued_successfully { + if let Some(err) = last_error { + return Err(err); + } else { + // NB: This branch should be impossible as we should always set last_error. + return Err(GrokError::Client(GrokClientError::ErrorGeneratingVideo)) + } + } + + let mut video_liked_successfully = false; + let mut last_error = None; + + for i in 0..3 { + info!("Liking media ... attempt {}", (i+1)); + + let request = GrokLikeMediaPost { + file_id: &upload_file_id, + cookie: current_full_credentials_ref.cookies.as_str(), + request_timeout: args.individual_request_timeout, + baggage: ¤t_full_credentials_ref.client_secrets.baggage, + sentry_trace: ¤t_full_credentials_ref.client_secrets.sentry_trace, + verification_token: ¤t_full_credentials_ref.client_secrets.verification_token, + svg_data: ¤t_full_credentials_ref.client_secrets.svg_path_data, + numbers: ¤t_full_credentials_ref.client_secrets.numbers, + }; + + let like_result = request.send().await; + + match like_result { + Ok(_) => { + info!("Media Liked"); + video_liked_successfully = true; + break; + } + Err(GrokError::ApiSpecific(GrokSpecificApiError::AutomationBlocked)) => { + info!("Grok automation blocked liking media; renewing credentials."); + last_error = Some(GrokError::ApiSpecific(GrokSpecificApiError::AutomationBlocked)); + } + Err(err) => { + info!("Grok error liking media; renewing credentials. Error = {:?}", err); + last_error = Some(err); + } + } + + info!("Refreshing Grok client secrets..."); + + let secrets = request_client_secrets(RequestClientSecretsArgs { + cookies: ¤t_full_credentials_ref.cookies, + }).await?; + + // NB: This is just to appease the borrow checker. It doesn't like that we're borrowing `maybe_new_credentials` + // in a loop and replacing it while borrowed, hence this hack. + current_full_credentials_ref = args.credentials; + + maybe_new_credentials = Some(GrokFullCredentials::from_cookies_and_client_secrets( + current_full_credentials_ref.cookies.clone(), + secrets, + )); + + if let Some(creds) = maybe_new_credentials.as_ref() { + current_full_credentials_ref = creds; + } + } + + if !video_liked_successfully { + if let Some(err) = last_error { + return Err(err); + } else { + // NB: This branch should be impossible as we should always set last_error. + return Err(GrokError::Client(GrokClientError::ErrorGeneratingVideo)) + } + } + + let maybe_video_url = maybe_video_file_id + .as_ref() + .map(|file_id| { + user_and_file_id_to_video_url(¤t_full_credentials_ref.client_secrets.user_id, file_id, false) + }); + + Ok(ImageUploadAndGenerateVideoWithRetryResult { + maybe_new_client_secrets: maybe_new_credentials + .map(|secrets| secrets.client_secrets), + upload_result: ImageUploadAndGenerateVideoResult { + post_id: post_result.post_id, + image_file_id: upload_file_id, + video_file_id: maybe_video_file_id, + video_url: maybe_video_url, + generation_is_complete, + } + }) +} + + +#[cfg(test)] +mod tests { + use crate::credentials::grok_full_credentials::GrokFullCredentials; + use crate::datatypes::api::aspect_ratio::AspectRatio; + use crate::datatypes::file_upload_spec::FileUploadSpec; + use crate::recipes::request_client_secrets::{request_client_secrets, RequestClientSecretsArgs}; + use crate::recipes::upload_image_and_generate_video_with_retry::{upload_image_and_generate_video_with_retry, UploadImageAndGenerateVideoWithRetry}; + use crate::test_utils::get_test_cookies::get_typed_test_cookies; + use crate::test_utils::setup_test_logging::setup_test_logging; + use errors::AnyhowResult; + use log::LevelFilter; + + // Result: GrokUploadFileResponse { file_metadata_id: + // Some("acdee48f-9d6f-4bc6-9d06-fcc97dd4418a"), file_uri: + // Some("users/85980643-ffab-4984-a3de-59a608c47d7f/acdee48f-9d6f-4bc6-9d06-fcc97dd4418a/content") } + + #[tokio::test] + #[ignore] // Client side tests only + async fn test_upload_image_and_generate_video() -> AnyhowResult<()> { + setup_test_logging(LevelFilter::Info); + + //let image_path = "/Users/bt/Pictures/Zelda 64 Art/FCgYX6tWEAEhpsy.jpg"; + //let prompt = "our hero link plunges the sword into the pedestal, the temple is glowing with a blue aura"; + + let image_path = "/Users/bt/Pictures/Midjourney/hipster_girl.png"; + let maybe_prompt = Some("an anime girl turns and high fives a giant t-rex"); + + let cookies = get_typed_test_cookies()?; + + let mut bad_secrets = request_client_secrets(RequestClientSecretsArgs { + cookies: &cookies, + }).await?; + + // NB: We're purposely messing these up to trigger "retry". + bad_secrets.numbers.numbers = Vec::new(); + bad_secrets.sentry_trace.0 = "".to_string(); + + let credentials = GrokFullCredentials::from_cookies_and_client_secrets(cookies, bad_secrets); + + let result = upload_image_and_generate_video_with_retry(UploadImageAndGenerateVideoWithRetry { + credentials: &credentials, + file: FileUploadSpec::Path(image_path), + prompt: maybe_prompt, + aspect_ratio: Some(AspectRatio::Square), + individual_request_timeout: None, + wait_for_generation: false, + }).await?; + + let new_secrets = result.maybe_new_client_secrets; + + println!("[test] New secrets ? : {:?}", new_secrets.is_some()); + + if let Some(secrets) = new_secrets { + println!("[test] Verification Token: {:?}", secrets.verification_token); + println!("[test] Sentry Trace: {:?}", secrets.sentry_trace); + println!("[test] Numbers: {:?}", secrets.numbers); + println!("[test] Svg Path: {:?}", secrets.svg_path_data); + println!("[test] Baggage: {:?}", secrets.baggage); + } + + let video_result = result.upload_result; + + println!("[test] Post ID: {:?}", video_result.post_id); + println!("[test] Image File ID: {:?}", video_result.image_file_id); + println!("[test] Video File ID: {:?}", video_result.video_file_id); + println!("[test] Video URL: {:?}", video_result.video_url); + + assert_eq!(1, 2); + + Ok(()) + } +} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs index f4321ef50..5ce8f3a27 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs @@ -23,6 +23,7 @@ use grok_client::datatypes::api::aspect_ratio::AspectRatio; use grok_client::datatypes::file_upload_spec::FileUploadSpec; use grok_client::recipes::request_client_secrets::{request_client_secrets, RequestClientSecretsArgs}; use grok_client::recipes::upload_image_and_generate_video::{upload_image_and_generate_video, UploadImageAndGenerateVideo}; +use grok_client::recipes::upload_image_and_generate_video_with_retry::{upload_image_and_generate_video_with_retry, UploadImageAndGenerateVideoWithRetry}; use log::{error, info, warn}; use openai_sora_client::recipes::generate_sora2_video::generate_sora2_video_with_session_auto_renew::generate_sora2_video_with_session_auto_renew; use openai_sora_client::requests::generate_sora2_video::generate_sora2_video::{GenerateSora2VideoArgs, Orientation}; @@ -65,8 +66,8 @@ pub async fn handle_grok_video( info!("Calling Grok Video generate..."); - let upload = upload_image_and_generate_video(UploadImageAndGenerateVideo { - full_credentials: &creds, + let upload_result = upload_image_and_generate_video_with_retry(UploadImageAndGenerateVideoWithRetry { + credentials: &creds, file: FileUploadSpec::Path(local_image_file.path()), prompt: request.prompt.as_deref(), aspect_ratio: Some(aspect_ratio), @@ -74,14 +75,18 @@ pub async fn handle_grok_video( individual_request_timeout: Some(GROK_IMAGE_UPLOAD_TIMEOUT) }).await; - let post_id = match upload { + let post_id = match upload_result { Err(err) => { error!("Failed to use Grok video: {:?}", err); return Err(GenerateError::from(err)); } Ok(result) => { - info!("Successfully enqueued Grok Video: post_id = {:?}", result.post_id); - result.post_id + if let Some(secrets) = result.maybe_new_client_secrets { + info!("New Grok secrets generated! Storing them in the credential manager..."); + grok_credential_manager.replace_client_secrets_only(secrets)?; + } + info!("Successfully enqueued Grok Video: post_id = {:?}", result.upload_result.post_id); + result.upload_result.post_id } }; diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs index d76267b12..286d7b07f 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs @@ -102,7 +102,7 @@ impl GrokCredentialManager { } } - pub fn replace_full_credentials(&self, creds: GrokFullCredentials) -> anyhow::Result<()> { + pub fn replace_full_credentials(&self, creds: GrokFullCredentials) -> anyhow::Result<()> { match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), Ok(mut holder) => { @@ -112,6 +112,22 @@ impl GrokCredentialManager { } } + pub fn replace_client_secrets_only(&self, secrets: GrokClientSecrets) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut holder) => { + let mut maybe_full_creds = holder.grok_full_credentials.clone(); + if let Some(creds) = maybe_full_creds.as_mut() { + creds.client_secrets = secrets; + } else { + warn!("No existing credentials to replace secrets into."); + } + holder.grok_full_credentials = maybe_full_creds; + Ok(()) + } + } + } + // NB: This is just a heuristic. We'll add better checks later. pub fn do_task_polling(&self) -> anyhow::Result { self.session_appears_active() diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs index 4c065d4b2..69f22880c 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs @@ -18,6 +18,7 @@ use enums::common::model_type::ModelType; use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass; use enums::tauri::tasks::task_model_type::TaskModelType; use errors::AnyhowResult; +use grok_client::recipes::prompt_websocket_image_with_retry::{prompt_websocket_image_with_retry, PromptWebsocketImageWithRetryArgs}; use grok_client::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; use grok_client::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; @@ -86,47 +87,33 @@ async fn inner_loop( continue; } - let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?; + let maybe_cookies = grok_creds.maybe_copy_typed_cookies()?; - let cookie_header = match maybe_cookie_header { - Some(cookie_header) => cookie_header, + //let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?; + + let cookies = match maybe_cookies { + Some(cookies) => cookies, None => { tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; continue; - } + }, }; + //let cookie_header = match maybe_cookie_header { + // Some(cookie_header) => cookie_header, + // None => { + // tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + // continue; + // } + //}; + let websocket = create_listen_websocket(CreateListenWebsocketArgs { - cookies: &cookie_header, + cookies: cookies.as_str(), }).await?; let mut websocket = GrokWebsocket::new(websocket); - //poll_task_loop( - // &app_handle, - // &app_env_configs, - // &app_data_root, - // &task_database, - // &grok_creds, - // websocket, - // prompt_queue, - // &storyteller_creds_manager, - //).await?; - loop { - // let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs { - // db: task_database.get_connection(), - // provider: GenerationProvider::Grok, - // task_statuses: &TASK_DATABASE_PENDING_STATUSES, - // }).await?; - // // Only Grok images - // let local_tasks = local_tasks.tasks.into_iter() - // .filter(|task| match task.model_type { - // Some(TaskModelType::GrokImage) => true, - // _ => false, - // }) - // .collect::>(); - let maybe_prompt = prompt_queue.dequeue()?; let prompt_item = match maybe_prompt { @@ -137,14 +124,20 @@ async fn inner_loop( } }; - info!("Prompting Grok websocket: {}", prompt_item.prompt); + info!("Prompt received over in-memory queue. Prompting Grok websocket: {}", prompt_item.prompt); - let _result = prompt_websocket_image(PromptWebsocketImageArgs { + let result = prompt_websocket_image_with_retry(PromptWebsocketImageWithRetryArgs { websocket: &mut websocket, prompt: &prompt_item.prompt, aspect_ratio: prompt_item.aspect_ratio, + cookies: &cookies, }).await?; + if let Some(new_websocket) = result.maybe_new_websocket { + info!("Replacing websocket with new one..."); + websocket = new_websocket; + } + let images = listen_for_websocket_images(ListenForWebsocketImagesArgs { websocket: &mut websocket, timeout: Duration::from_millis(30_000), @@ -164,19 +157,6 @@ async fn inner_loop( } } -//async fn poll_task_loop( -// app_handle: &AppHandle, -// app_env_configs: &AppEnvConfigs, -// app_data_root: &AppDataRoot, -// task_database: &TaskDatabase, -// grok_creds: &GrokCredentialManager, -// grok_websocket: GrokWebsocket, -// prompt_queue: &GrokImagePromptQueue, -// storyteller_creds_manager: &StorytellerCredentialManager, -//) -> AnyhowResult<()> { -// -//} - async fn upload_images_to_storyteller( app_handle: &AppHandle, app_env_configs: &AppEnvConfigs,