hm, sync/send issues with a mutex/arc?

crates/api_clients/grok_client/src/datatypes/api/request_id.rs

@@ -5,3 +5,9 @@
 /// Request IDs are UUIds. 
 #[derive(Clone, Debug)]
 pub struct RequestId(pub String);
+
+impl RequestId {
+  pub fn to_string(&self) -> String {
+    self.0.clone()
+  }
+}

crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs

@@ -143,7 +143,7 @@ mod tests {
     std::io::stdout().flush()?;
 
     let _result = prompt_websocket_image(PromptWebsocketImageArgs {
-      websocket_wrapped: websocket.clone(),
+      websocket_wrapped: &websocket,
       prompt,
     }).await?;
 

crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs

@@ -100,7 +100,7 @@ mod tests {
     std::io::stdout().flush()?;
 
     let _result = prompt_websocket_image(PromptWebsocketImageArgs {
-      websocket_wrapped: websocket.clone(),
+      websocket_wrapped: &websocket,
       prompt,
     }).await?;
 

crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs

@@ -3,7 +3,7 @@ use crate::requests::image_websocket::grok_websocket::GrokWebsocket;
 use crate::requests::image_websocket::messages::websocket_client_message::WebsocketClientMessage;
 
 pub struct PromptWebsocketImageArgs<'a> {
-  pub websocket_wrapped: GrokWebsocket,
+  pub websocket_wrapped: &'a GrokWebsocket,
   pub prompt: &'a str,
 }
 
@@ -46,7 +46,7 @@ mod tests {
     std::io::stdout().flush()?;
 
     let result = prompt_websocket_image(PromptWebsocketImageArgs {
-      websocket_wrapped: websocket.clone(),
+      websocket_wrapped: &websocket,
       prompt: "a dog riding a motorcycle",
     }).await?;
 

crates/desktop/artcraft/src/core/artcraft_error.rs

@@ -11,6 +11,9 @@ pub enum ArtcraftError {
   // Service errors
   GrokError(GrokError),
   StorytellerError(StorytellerError),
+  // Lock errors
+  RwLockReadError,
+  RwLockWriteError,
 }
 
 impl Error for ArtcraftError {}
@@ -23,6 +26,8 @@ impl Display for ArtcraftError {
       Self::IoError(e) => write!(f, "IoError: {:?}", e),
       Self::GrokError(e) => write!(f, "GrokError: {:?}", e),
       Self::StorytellerError(e) => write!(f, "StorytellerError: {:?}", e),
+      Self::RwLockReadError => write!(f, "RwLockReadError"),
+      Self::RwLockWriteError => write!(f, "RwLockWriteError"),
     }
   }
 }

crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs

@@ -1,4 +1,5 @@
 use crate::core::artcraft_error::ArtcraftError;
+use anyhow::anyhow;
 use base64::DecodeError;
 use errors::AnyhowError;
 use grok_client::error::grok_error::GrokError;
@@ -79,6 +80,8 @@ pub enum BillingProvider {
 #[derive(Debug)]
 pub enum ProviderFailureReason {
   GrokError(GrokError),
+  /// NB: The Grok client doesn't say why certain errors (eg. missing fields) happen, so we synthesize this.
+  GrokJobEnqueueFailed,
   //Fal(FalErrorPlus),
   MidjourneyError(MidjourneyError),
   /// NB: The midjourney client doesn't categorize all errors, so we have to do so on our end.
@@ -148,6 +151,8 @@ impl From<ArtcraftError> for GenerateError {
       ArtcraftError::IoError(e) => Self::IoError(e),
       ArtcraftError::GrokError(e) => Self::ProviderFailure(ProviderFailureReason::GrokError(e)),
       ArtcraftError::StorytellerError(e) => Self::ProviderFailure(ProviderFailureReason::StorytellerError(e)),
+      ArtcraftError::RwLockReadError => Self::AnyhowError(anyhow!("Lock read error")),
+      ArtcraftError::RwLockWriteError => Self::AnyhowError(anyhow!("Lock write error")),
     }
   }
 }

crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs

@@ -67,6 +67,7 @@ impl TaskEnqueueSuccess{
       Some(GenerationModel::FluxProKontextMax) => Some(TaskModelType::FluxProKontextMax),
       Some(GenerationModel::Gemini25Flash) => Some(TaskModelType::Gemini25Flash),
       Some(GenerationModel::GptImage1) => Some(TaskModelType::GptImage1),
+      Some(GenerationModel::GrokImage) => Some(TaskModelType::GrokImage),
       Some(GenerationModel::Recraft3) => Some(TaskModelType::Recraft3),
       Some(GenerationModel::GrokVideo) => Some(TaskModelType::GrokVideo),
       Some(GenerationModel::Kling21Pro) => Some(TaskModelType::Kling21Pro),

crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs

@@ -6,6 +6,7 @@ use crate::core::commands::enqueue::text_to_image::generic::handle_image_artcraf
 use crate::core::commands::enqueue::text_to_image::generic::handle_image_fal::handle_image_fal;
 use crate::core::commands::enqueue::text_to_image::gpt_image_1::handle_gpt_image_1::handle_gpt_image_1;
 use crate::core::commands::enqueue::text_to_image::gpt_image_1::handle_gpt_image_1_sora::handle_gpt_image_1_sora;
+use crate::core::commands::enqueue::text_to_image::grok::handle_grok::handle_grok;
 use crate::core::commands::enqueue::text_to_image::midjourney::handle_midjourney::handle_midjourney;
 use crate::core::commands::response::failure_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus};
 use crate::core::commands::response::shorthand::Response;
@@ -19,6 +20,8 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::provider_priority::{Provider, ProviderPriorityStore};
 use crate::core::state::task_database::TaskDatabase;
+use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
+use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
 use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
@@ -48,6 +51,8 @@ pub enum ImageModel {
   FluxPro11Ultra,
   #[serde(rename = "gpt_image_1")]
   GptImage1,
+  #[serde(rename = "grok_image")]
+  GrokImage,
   #[serde(rename = "recraft_3")]
   Recraft3,
 
@@ -144,6 +149,8 @@ pub async fn enqueue_text_to_image_command(
   provider_priority_store: State<'_, ProviderPriorityStore>,
   task_database: State<'_, TaskDatabase>,
   mj_creds_manager: State<'_, MidjourneyCredentialManager>,
+  grok_creds_manager: State<'_, GrokCredentialManager>,
+  grok_websocket_manager: State<'_, GrokWebsocketManager>,
   storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
   sora_creds_manager: State<'_, SoraCredentialManager>,
   sora_task_queue: State<'_, SoraTaskQueue>,
@@ -158,6 +165,8 @@ pub async fn enqueue_text_to_image_command(
     &provider_priority_store,
     &task_database,
     &mj_creds_manager,
+    &grok_creds_manager,
+    &grok_websocket_manager,
     &storyteller_creds_manager,
     &app_env_configs,
     &sora_creds_manager,
@@ -226,6 +235,8 @@ pub async fn handle_request(
   provider_priority_store: &ProviderPriorityStore,
   task_database: &TaskDatabase,
   mj_creds_manager: &MidjourneyCredentialManager,
+  grok_creds_manager: &GrokCredentialManager,
+  grok_websocket_manager: &GrokWebsocketManager,
   storyteller_creds_manager: &StorytellerCredentialManager,
   app_env_configs: &AppEnvConfigs,
   sora_creds_manager: &SoraCredentialManager,
@@ -240,6 +251,8 @@ pub async fn handle_request(
     &storyteller_creds_manager,
     &app_env_configs,
     &mj_creds_manager,
+    &grok_creds_manager,
+    &grok_websocket_manager,
     &sora_creds_manager,
     &sora_task_queue,
   ).await;
@@ -274,6 +287,8 @@ pub async fn dispatch_request(
   storyteller_creds_manager: &StorytellerCredentialManager,
   app_env_configs: &AppEnvConfigs,
   mj_creds_manager: &MidjourneyCredentialManager,
+  grok_creds_manager: &GrokCredentialManager,
+  grok_websocket_manager: &GrokWebsocketManager,
   sora_creds_manager: &SoraCredentialManager,
   sora_task_queue: &SoraTaskQueue,
 ) -> Result<TaskEnqueueSuccess, GenerateError> {
@@ -306,6 +321,15 @@ pub async fn dispatch_request(
         sora_task_queue,
       ).await;
     }
+    Some(ImageModel::GrokImage) => {
+      return handle_grok(
+        app,
+        request,
+        app_env_configs,
+        grok_creds_manager,
+        grok_websocket_manager,
+      )
+    }
     Some(ImageModel::Midjourney) => {
       return handle_midjourney(
         app,

crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs

@@ -0,0 +1,155 @@
+use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason, ProviderFailureReason};
+use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize};
+use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
+use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize};
+use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
+use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent;
+use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent;
+use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
+use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
+use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::state::provider_priority::ProviderPriorityStore;
+use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
+use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager;
+use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
+use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use artcraft_api_defs::generate::image::generate_gpt_image_1_text_to_image::{GenerateGptImage1TextToImageImageQuality, GenerateGptImage1TextToImageImageSize, GenerateGptImage1TextToImageNumImages, GenerateGptImage1TextToImageRequest};
+use enums::common::generation_provider::GenerationProvider;
+use enums::tauri::tasks::task_type::TaskType;
+use grok_client::error::grok_error::GrokError;
+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::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs};
+use grok_client::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs};
+use idempotency::uuid::generate_random_uuid;
+use log::{error, info, warn};
+use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
+use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest};
+use midjourney_client::error::midjourney_api_error::MidjourneyApiError;
+use midjourney_client::recipes::channel_id::ChannelId;
+use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageError, TextToImageRequest};
+use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
+use openai_sora_client::recipes::simple_image_gen_with_session_auto_renew::{simple_image_gen_with_session_auto_renew, SimpleImageGenAutoRenewRequest};
+use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages};
+use std::time::Duration;
+use storyteller_client::endpoints::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image;
+use storyteller_client::endpoints::generate::image::generate_gpt_image_1_text_to_image::generate_gpt_image_1_text_to_image;
+use tauri::AppHandle;
+use tokens::tokens::media_files::MediaFileToken;
+
+pub async fn handle_grok(
+  app: &AppHandle,
+  request: &EnqueueTextToImageRequest,
+  app_env_configs: &AppEnvConfigs,
+  creds_manager: &GrokCredentialManager,
+  websocket_manager: &GrokWebsocketManager,
+) -> Result<TaskEnqueueSuccess, GenerateError> {
+
+  // TODO: We can request population of the user info if absent or expired.
+  let websocket = get_websocket(
+    app,
+    creds_manager,
+    websocket_manager,
+  ).await?;
+  
+  let _result = prompt_websocket_image(PromptWebsocketImageArgs {
+    websocket_wrapped: &websocket,
+    prompt: request.prompt.as_deref().unwrap_or_else(||""),
+  }).await?;
+  
+  let request_id = listen_for_websocket_request_id(ListenForWebsocketRequestIdArgs {
+    websocket: &websocket,
+    timeout: Duration::from_millis(10_000),
+  }).await?;
+  
+  let request_id = match request_id.request_id {
+    Some(request_id) => request_id.to_string(),
+    None => {
+      warn!("No request id from Grok image generation (websocket)");
+      return Err(GenerateError::ProviderFailure(ProviderFailureReason::GrokJobEnqueueFailed));
+    }
+  };
+  
+  info!("Successfully enqueued MidJourney. Job token: {}", request_id);
+
+  Ok(TaskEnqueueSuccess {
+    provider: GenerationProvider::Grok,
+    model: Some(GenerationModel::GrokImage),
+    provider_job_id: Some(request_id),
+    task_type: TaskType::ImageGeneration,
+  })
+}
+
+// TODO:
+// fn handle_midjourney_errors(
+//   app: &AppHandle,
+//   maybe_errors: Option<Vec<TextToImageError>>
+// ) -> Result<TaskEnqueueSuccess, GenerateError> {
+//   if let Some(errors) = maybe_errors {
+//     if !errors.is_empty() {
+//       let messages: Vec<String> = errors.iter()
+//           .map(|e| format!("{:?}", e))
+//           .collect();
+// 
+//       let combined_message = messages.join("; ");
+// 
+//       let event = FlashUserInputErrorEvent {
+//         message: format!("Midjourney Error: {}", combined_message),
+//       };
+// 
+//       if let Err(err) = event.send(&app) {
+//         error!("Failed to send FlashUserInputErrorEvent: {:?}", err); // Fail open
+//       }
+//     }
+//   }
+//   
+//   Err(GenerateError::ProviderFailure(ProviderFailureReason::MidjourneyJobEnqueueFailed))
+// }
+
+async fn get_websocket(app: &AppHandle, creds_manager: &GrokCredentialManager, websocket_manager: &GrokWebsocketManager) -> Result<GrokWebsocket, GenerateError> {
+  let maybe_websocket = websocket_manager.grab_websocket()?;
+  
+  let websocket = match maybe_websocket {
+    Some(websocket) => websocket,
+    None => {
+      let cookies = match creds_manager.maybe_copy_cookie_store() {
+        Ok(Some(cookies)) => cookies.to_cookie_string(),
+        Ok(None) => {
+          error!("Grok credentials not found.");
+          ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Grok, &app);
+          return Err(GenerateError::needs_grok_credentials());
+        }
+        Err(err) => {
+          error!("Error reading Midjourney credentials: {:?}", err);
+          ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Grok, &app);
+          return Err(GenerateError::needs_grok_credentials());
+        },
+      };
+      
+      info!("Creating Grok websocket...");
+      
+      let create_websocket_result = create_listen_websocket(CreateListenWebsocketArgs {
+        cookies: &cookies,
+      }).await;
+      
+      let websocket = match create_websocket_result {
+        Ok(websocket) => GrokWebsocket::new(websocket),
+        Err(err) => {
+          error!("Error creating websocket: {:?}", err);
+          ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Grok, &app);
+          return Err(GenerateError::needs_grok_credentials());
+        }
+      };
+      
+      info!("Setting managed Grok websocket...");
+      websocket_manager.set_websocket(websocket.clone())?;
+      
+      websocket
+    }
+  };
+  
+  Ok(websocket)
+}

crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs

@@ -0,0 +1 @@
+pub (super) mod handle_grok;

crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs

@@ -2,4 +2,5 @@ pub mod enqueue_text_to_image_command;
 pub mod generic;
 mod gemini_25_flash;
 mod gpt_image_1;
-mod midjourney;
+mod midjourney;
+mod grok;

crates/desktop/artcraft/src/core/events/generation_events/common.rs

@@ -30,6 +30,8 @@ pub enum GenerationModel {
   FluxProKontextMax,
   #[serde(rename = "gpt_image_1")]
   GptImage1,
+  #[serde(rename = "grok_image")]
+  GrokImage,
   #[serde(rename = "gemini_25_flash")]
   Gemini25Flash,
   

crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs

@@ -10,6 +10,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::artcraft_platform_info::ArtcraftPlatformInfo;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
+use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager;
 use crate::services::grok::threads::grok_video_task_polling::grok_video_task_polling_thread::grok_video_task_polling_thread;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
 use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread;
@@ -29,6 +30,7 @@ pub async fn handle_tauri_startup(
   sora_task_queue: SoraTaskQueue,
   mj_creds_manager: MidjourneyCredentialManager,
   grok_creds_manager: GrokCredentialManager,
+  grok_websocket_manager: GrokWebsocketManager,
 ) -> AnyhowResult<()> {
 
   set_app_log_level(

crates/desktop/artcraft/src/lib.rs

@@ -33,6 +33,7 @@ use crate::services::grok::commands::grok_clear_credentials_command::grok_clear_
 use crate::services::grok::commands::grok_get_credential_info_command::grok_get_credential_info_command;
 use crate::services::grok::commands::grok_open_login_command::grok_open_login_command;
 use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
+use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager;
 use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command;
 use crate::services::midjourney::commands::midjourney_get_credential_info_command::midjourney_get_credential_info_command;
 use crate::services::midjourney::commands::midjourney_open_login_command::midjourney_open_login_command;
@@ -104,6 +105,9 @@ pub fn run() {
   let grok_creds_manager = GrokCredentialManager::initialize_from_disk_infallible(&app_data_root);
   let grok_creds_manager_2 = grok_creds_manager.clone();
 
+  let grok_websocket_manager = GrokWebsocketManager::new();
+  let grok_websocket_manager_2= grok_websocket_manager.clone();
+
   println!("Initializing backend runtime...");
 
   let builder = tauri::Builder::default()
@@ -142,6 +146,7 @@ pub fn run() {
           sora_tasks,
           midjourney_creds_manager_2,
           grok_creds_manager_2,
+          grok_websocket_manager_2,
         ).await;
 
         if let Err(err) = result {
@@ -157,6 +162,7 @@ pub fn run() {
     .manage(app_preferences)
     .manage(artcraft_platform_info)
     .manage(grok_creds_manager)
+    .manage(grok_websocket_manager)
     .manage(midjourney_creds_manager)
     .manage(sora_creds_manager)
     .manage(sora_task_queue)

crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs

@@ -1,7 +1,59 @@
+use crate::core::artcraft_error::ArtcraftError;
+use errors::AnyhowResult;
 use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket;
-use std::sync::{Arc, RwLock};
+use log::error;
+use std::sync::{Arc, LockResult, RwLock};
 
+
+/// NB: This is inefficient because the websockets are locked at two layers.
+/// Should be fine for our performance needs, though.
 #[derive(Clone)]
-pub struct GrokWebsocketHolder {
-  websocket: Arc<RwLock<GrokWebsocket>>,
-}
+pub struct GrokWebsocketManager {
+  websocket: Arc<RwLock<Option<GrokWebsocket>>>,
+}
+
+impl GrokWebsocketManager {
+  pub fn new() -> Self {
+    Self {
+      websocket: Arc::new(RwLock::new(None)),
+    }
+  }
+
+  pub fn set_websocket(&self, websocket: GrokWebsocket) -> Result<(), ArtcraftError> {
+    match self.websocket.write() {
+      Ok(mut guard) => {
+        *guard = Some(websocket);
+        Ok(())
+      }
+      Err(err) => {
+        error!("Error writing locked websocket: {}", err);
+        Err(ArtcraftError::RwLockWriteError)
+      }
+    }
+  }
+
+  pub fn clear_websocket(&self) -> Result<(), ArtcraftError> {
+    match self.websocket.write() {
+      Ok(mut guard) => {
+        *guard = None;
+        Ok(())
+      }
+      Err(err) => {
+        error!("Error writing locked websocket: {}", err);
+        Err(ArtcraftError::RwLockWriteError)
+      }
+    }
+  }
+
+  pub fn grab_websocket(&self) -> Result<Option<GrokWebsocket>, ArtcraftError> {
+    match self.websocket.read() {
+      Ok(guard) => {
+        Ok(guard.clone())
+      }
+      Err(err) => {
+        error!("Error reading locked websocket: {}", err);
+        Err(ArtcraftError::RwLockReadError)
+      }
+    }
+  }
+}

crates/desktop/artcraft/src/services/grok/state/mod.rs

@@ -1,4 +1,6 @@
 pub mod grok_credential_manager;
 pub mod grok_serializable_state;
+pub mod grok_websocket_manager;
+
+// Private
 mod grok_credential_holder;
-mod grok_websocket_manager;

crates/schema/public/enums/src/tauri/tasks/task_model_type.rs

@@ -31,6 +31,8 @@ pub enum TaskModelType {
   Gemini25Flash,
   #[serde(rename = "gpt_image_1")]
   GptImage1,
+  #[serde(rename = "grok_image")]
+  GrokImage,
   #[serde(rename = "recraft_3")]
   Recraft3,
   
@@ -84,6 +86,7 @@ impl TaskModelType {
       Self::FluxProKontextMax => "flux_pro_kontext_max",
       Self::Gemini25Flash => "gemini_25_flash",
       Self::GptImage1 => "gpt_image_1",
+      Self::GrokImage => "grok_image",
       Self::Recraft3 => "recraft_3",
       Self::Midjourney => "midjourney",
       // Video models
@@ -114,6 +117,7 @@ impl TaskModelType {
       "flux_pro_kontext_max" => Ok(Self::FluxProKontextMax),
       "gemini_25_flash" => Ok(Self::Gemini25Flash),
       "gpt_image_1" => Ok(Self::GptImage1),
+      "grok_image" => Ok(Self::GrokImage),
       "recraft_3" => Ok(Self::Recraft3),
       "midjourney" => Ok(Self::Midjourney),
       // Video models
@@ -147,6 +151,7 @@ impl TaskModelType {
       Self::FluxProKontextMax,
       Self::Gemini25Flash,
       Self::GptImage1,
+      Self::GrokImage,
       Self::Recraft3,
       Self::Midjourney,
       // Video models
@@ -187,6 +192,7 @@ mod tests {
       assert_serialization(TaskModelType::FluxProKontextMax, "flux_pro_kontext_max");
       assert_serialization(TaskModelType::Gemini25Flash, "gemini_25_flash");
       assert_serialization(TaskModelType::GptImage1, "gpt_image_1");
+      assert_serialization(TaskModelType::GrokImage, "grok_image");
       assert_serialization(TaskModelType::Recraft3, "recraft_3");
       assert_serialization(TaskModelType::Midjourney, "midjourney");
       // Video models
@@ -216,6 +222,7 @@ mod tests {
       assert_eq!(TaskModelType::FluxProKontextMax.to_str(), "flux_pro_kontext_max");
       assert_eq!(TaskModelType::Gemini25Flash.to_str(), "gemini_25_flash");
       assert_eq!(TaskModelType::GptImage1.to_str(), "gpt_image_1");
+      assert_eq!(TaskModelType::GrokImage.to_str(), "grok_image");
       assert_eq!(TaskModelType::Recraft3.to_str(), "recraft_3");
       assert_eq!(TaskModelType::Midjourney.to_str(), "midjourney");
       // Video models
@@ -245,6 +252,7 @@ mod tests {
       assert_eq!(TaskModelType::from_str("flux_pro_kontext_max").unwrap(), TaskModelType::FluxProKontextMax);
       assert_eq!(TaskModelType::from_str("gemini_25_flash").unwrap(), TaskModelType::Gemini25Flash);
       assert_eq!(TaskModelType::from_str("gpt_image_1").unwrap(), TaskModelType::GptImage1);
+      assert_eq!(TaskModelType::from_str("grok_image").unwrap(), TaskModelType::GrokImage);
       assert_eq!(TaskModelType::from_str("recraft_3").unwrap(), TaskModelType::Recraft3);
       assert_eq!(TaskModelType::from_str("midjourney").unwrap(), TaskModelType::Midjourney);
       // Video models
@@ -276,7 +284,7 @@ mod tests {
     #[test]
     fn all_variants() {
       let mut variants = TaskModelType::all_variants();
-      assert_eq!(variants.len(), 22);
+      assert_eq!(variants.len(), 23);
       // Image models
       assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev));
       assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell));
@@ -287,6 +295,7 @@ mod tests {
       assert_eq!(variants.pop_first(), Some(TaskModelType::FluxProKontextMax));
       assert_eq!(variants.pop_first(), Some(TaskModelType::Gemini25Flash));
       assert_eq!(variants.pop_first(), Some(TaskModelType::GptImage1));
+      assert_eq!(variants.pop_first(), Some(TaskModelType::GrokImage));
       assert_eq!(variants.pop_first(), Some(TaskModelType::Recraft3));
       assert_eq!(variants.pop_first(), Some(TaskModelType::Midjourney));
       // Video models