grok retry (#916)

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."),

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;

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<GrokWebsocket>,
+}
+
+
+pub async fn prompt_websocket_image_with_retry(args: PromptWebsocketImageWithRetryArgs<'_>) -> Result<PromptWebsocketImageWithRetryResult, GrokError>
+{
+  let mut maybe_new_websocket : Option<GrokWebsocket> = 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<String> = 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(())
+  }
+}

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<Path>> {
   pub credentials: &'a GrokFullCredentials,
@@ -12,6 +27,9 @@ pub struct UploadImageAndGenerateVideoWithRetry<'a, P: AsRef<Path>> {
   /// Video generation prompt
   pub prompt: Option<&'a str>,
 
+  /// Aspect ratio for the video.
+  pub aspect_ratio: Option<AspectRatio>,
+
   /// 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<Path>> {
 
   pub individual_request_timeout: Option<Duration>,
 }
+
+#[derive(Clone)]
+pub struct ImageUploadAndGenerateVideoWithRetryResult {
+  /// If new secrets were generated, they're stored here.
+  pub maybe_new_client_secrets: Option<GrokClientSecrets>,
+
+  /// The video generation result.
+  pub upload_result: ImageUploadAndGenerateVideoResult,
+}
+
+pub async fn upload_image_and_generate_video_with_retry<P: AsRef<Path>>(
+  args: UploadImageAndGenerateVideoWithRetry<'_, P>
+) -> Result<ImageUploadAndGenerateVideoWithRetryResult, GrokError> {
+
+  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(
+    &current_full_credentials_ref.client_secrets.user_id,
+    &upload_file_id
+  );
+
+  info!("Uploaded URI: {:?}", url);
+
+  info!("Creating media post...");
+
+  let request = GrokCreateMediaPost {
+    user_id: &current_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: &current_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: &current_full_credentials_ref.client_secrets.baggage,
+      sentry_trace: &current_full_credentials_ref.client_secrets.sentry_trace,
+      verification_token: &current_full_credentials_ref.client_secrets.verification_token,
+      svg_data: &current_full_credentials_ref.client_secrets.svg_path_data,
+      numbers: &current_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: &current_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: &current_full_credentials_ref.client_secrets.baggage,
+      sentry_trace: &current_full_credentials_ref.client_secrets.sentry_trace,
+      verification_token: &current_full_credentials_ref.client_secrets.verification_token,
+      svg_data: &current_full_credentials_ref.client_secrets.svg_path_data,
+      numbers: &current_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: &current_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(&current_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(())
+  }
+}

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

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<bool> {
     self.session_appears_active()

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::<Vec<_>>();
-
       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,