grok images (#893)

.env

@@ -1,6 +1,7 @@
 SERVER_ENVIRONMENT=development
 
 DATABASE_URL=mysql://storyteller:password@localhost/storyteller
+#DATABASE_URL=sqlite:/tmp/tasks.sqlite
 #DATABASE_URL=mysql://root:root@localhost/storyteller
 #DATABASE_URL=mysql://root@localhost/storyteller
 

.sqlx/query-1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68.json

@@ -0,0 +1,12 @@
+{
+  "db_name": "SQLite",
+  "query": "\n    UPDATE tasks\n    SET\n      task_status = ?,\n      on_complete_batch_token = ?,\n      on_complete_primary_media_file_token = ?,\n      on_complete_primary_media_file_class = ?,\n      on_complete_primary_media_file_cdn_url = ?,\n      on_complete_primary_media_file_thumbnail_url_template = ?\n    WHERE\n      provider = ?\n      AND provider_job_id = ?\n  ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Right": 8
+    },
+    "nullable": []
+  },
+  "hash": "1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68"
+}

.sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json

@@ -0,0 +1,68 @@
+{
+  "db_name": "SQLite",
+  "query": "\n    SELECT\n      id,\n      task_status,\n      task_type,\n      model_type,\n      provider,\n      provider_job_id,\n      frontend_caller,\n      frontend_subscriber_id,\n      frontend_subscriber_payload\n    FROM tasks\n    WHERE\n      provider = ?\n      AND provider_job_id = ?\n  ",
+  "describe": {
+    "columns": [
+      {
+        "name": "id",
+        "ordinal": 0,
+        "type_info": "Text"
+      },
+      {
+        "name": "task_status",
+        "ordinal": 1,
+        "type_info": "Text"
+      },
+      {
+        "name": "task_type",
+        "ordinal": 2,
+        "type_info": "Text"
+      },
+      {
+        "name": "model_type",
+        "ordinal": 3,
+        "type_info": "Text"
+      },
+      {
+        "name": "provider",
+        "ordinal": 4,
+        "type_info": "Text"
+      },
+      {
+        "name": "provider_job_id",
+        "ordinal": 5,
+        "type_info": "Text"
+      },
+      {
+        "name": "frontend_caller",
+        "ordinal": 6,
+        "type_info": "Text"
+      },
+      {
+        "name": "frontend_subscriber_id",
+        "ordinal": 7,
+        "type_info": "Text"
+      },
+      {
+        "name": "frontend_subscriber_payload",
+        "ordinal": 8,
+        "type_info": "Text"
+      }
+    ],
+    "parameters": {
+      "Right": 2
+    },
+    "nullable": [
+      false,
+      false,
+      false,
+      true,
+      true,
+      true,
+      true,
+      true,
+      true
+    ]
+  },
+  "hash": "38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069"
+}

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/grok_websocket.rs

@@ -0,0 +1,108 @@
+use crate::error::grok_client_error::GrokClientError;
+use crate::error::grok_error::GrokError;
+use crate::requests::image_websocket::messages::websocket_server_message::WebsocketServerMessage;
+use futures::stream::FusedStream;
+use futures::{Stream, TryStreamExt};
+use log::{info, warn};
+use serde::Serialize;
+use std::sync::{Arc, RwLock};
+use std::time::Duration;
+use wreq::ws::message::Message;
+use wreq::ws::WebSocket;
+
+pub struct GrokWebsocket {
+  pub(crate) websocket: WebSocket,
+}
+
+impl GrokWebsocket {
+  pub fn new(websocket: WebSocket) -> Self {
+    Self {
+      websocket,
+    }
+  }
+
+  pub async fn send(&mut self, message: String) -> Result<(), GrokError> {
+    let message = Message::text(message);
+    self.websocket.send(message)
+        .await
+        .map_err(|err| {
+          GrokClientError::WreqClientError(err)
+        })?;
+    Ok(())
+  }
+
+  pub async fn send_serializable<T: Serialize>(&mut self, message: T) -> Result<(), GrokError> {
+    let message_json = serde_json::to_string(&message)
+        .map_err(|err| {
+          warn!("Failed to serialize prompt websocket message: {}", err);
+          GrokClientError::WebsocketRequestSerializationError(err)
+        })?;
+    self.send(message_json).await
+  }
+
+  pub async fn get_response_with_timeout(&mut self, duration: Duration) -> Result<Option<WebsocketServerMessage>, GrokError> {
+    let result = tokio::time::timeout(duration, self.websocket.try_next()).await;
+
+    match result {
+      Err(elapsed) => {
+        info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);
+        Ok(None) // Timeout elapsed.
+      }
+      Ok(inner) => {
+        let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;
+
+        match maybe_message {
+          None => Ok(None),
+          Some(Message::Text(text)) => {
+            let maybe_message = WebsocketServerMessage::from_json_str(&text)?;
+            Ok(Some(maybe_message))
+          },
+          Some(_) => {
+            warn!("Received non-text websocket message.");
+            Ok(None)
+          },
+        }
+      }
+    }
+  }
+
+  pub async fn try_next_timeout(&mut self, duration: Duration) -> Result<Option<Message>, GrokClientError> {
+    let result = tokio::time::timeout(duration, self.websocket.try_next()).await;
+
+    match result {
+      Err(elapsed) => {
+        info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);
+        Ok(None) // Timeout elapsed.
+      }
+      Ok(inner) => {
+        let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;
+        Ok(maybe_message)
+      }
+    }
+  }
+
+  // pub async fn is_terminated(&self) -> Result<bool, GrokClientError> {
+  //   match self.websocket.read() {
+  //     Err(_) => Err(GrokClientError::WebsocketLockError.into()),
+  //     Ok(websocket) => Ok(websocket.is_terminated()),
+  //   }
+  // }
+
+  // /// NB: This will block forever if there is no item.
+  // pub async fn try_next(&self) -> Result<Option<Message>, GrokClientError> {
+  //   match self.websocket.write() {
+  //     Err(_) => Err(GrokClientError::WebsocketLockError),
+  //     Ok(mut websocket) => websocket
+  //         .try_next()
+  //         .await
+  //         .map_err(|err| GrokClientError::WebsocketReadError(err)),
+  //   }
+  // }
+
+  // pub fn size_hint(&self) -> Result<(usize, Option<usize>), GrokClientError> {
+  //   match self.websocket.read() {
+  //     Err(_) => Err(GrokClientError::WebsocketLockError),
+  //     Ok(websocket) => Ok(websocket.size_hint())
+  //   }
+  // }
+}

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

@@ -11,11 +11,11 @@ use wreq::ws::message::Message;
 use wreq::ws::WebSocket;
 
 #[derive(Clone)]
-pub struct ClonableWebsocket {
+pub struct GrokWrappedWebsocket {
   pub(crate) websocket: Arc<RwLock<WebSocket>>,
 }
 
-impl ClonableWebsocket {
+impl GrokWrappedWebsocket {
   pub fn new(websocket: WebSocket) -> Self {
     Self {
       websocket: Arc::new(RwLock::new(websocket)),
@@ -77,7 +77,6 @@ impl ClonableWebsocket {
     }
   }
 
-
   pub async fn try_next_timeout(&self, duration: Duration) -> Result<Option<Message>, GrokClientError> {
     match self.websocket.write() {
       Err(_) => Err(GrokClientError::WebsocketLockError),

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

@@ -1,7 +1,8 @@
 use crate::datatypes::api::request_id::RequestId;
 use crate::error::grok_client_error::GrokClientError;
 use crate::error::grok_error::GrokError;
-use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
+use crate::requests::image_websocket::grok_websocket::GrokWebsocket;
+use crate::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket;
 use crate::requests::image_websocket::messages::message_image_data::MessageImageData;
 use crate::utils::scrub_blobs_for_debug_logging::scrub_blobs_for_debug_logging;
 use log::info;
@@ -12,7 +13,7 @@ use wreq::ws::message::Message;
 const DEFAULT_IMAGE_COUNT: usize = 6;
 
 pub struct ListenForWebsocketImagesArgs<'a> {
-  pub websocket: &'a ClonableWebsocket,
+  pub websocket: &'a mut GrokWebsocket,
   pub timeout: Duration,
 }
 
@@ -113,8 +114,8 @@ pub async fn listen_for_websocket_images(args: ListenForWebsocketImagesArgs<'_>)
 
 #[cfg(test)]
 mod tests {
-  use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
   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::listen_for_websocket_images::{listen_for_websocket_images, ListenForWebsocketImagesArgs};
   use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs};
   use crate::test_utils::get_test_cookies::get_test_cookies;
@@ -137,13 +138,14 @@ mod tests {
       cookies: &cookies,
     }).await?;
 
-    let websocket = ClonableWebsocket::new(websocket);
+    //let websocket = GrokWrappedWebsocket::new(websocket);
+    let mut websocket = GrokWebsocket::new(websocket);
 
     println!("Sending...");
     std::io::stdout().flush()?;
 
     let _result = prompt_websocket_image(PromptWebsocketImageArgs {
-      websocket_wrapped: websocket.clone(),
+      websocket: &mut websocket,
       prompt,
     }).await?;
 
@@ -153,7 +155,7 @@ mod tests {
     println!("Polling.");
 
     let images = listen_for_websocket_images(ListenForWebsocketImagesArgs {
-      websocket: &websocket,
+      websocket: &mut websocket,
       timeout: Duration::from_millis(10_000),
     }).await?;
 

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

@@ -1,13 +1,13 @@
 use crate::datatypes::api::request_id::RequestId;
 use crate::error::grok_client_error::GrokClientError;
 use crate::error::grok_error::GrokError;
-use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
+use crate::requests::image_websocket::grok_websocket::GrokWebsocket;
 use crate::requests::image_websocket::messages::message_request_id::MessageRequestId;
 use std::time::{Duration, Instant};
 use wreq::ws::message::Message;
 
 pub struct ListenForWebsocketRequestIdArgs<'a> {
-  pub websocket: &'a ClonableWebsocket,
+  pub websocket: &'a mut GrokWebsocket,
   pub timeout: Duration,
 }
 
@@ -70,8 +70,8 @@ pub async fn listen_for_websocket_request_id(args: ListenForWebsocketRequestIdAr
 
 #[cfg(test)]
 mod tests {
-  use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
   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::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs};
   use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs};
   use crate::test_utils::get_test_cookies::get_test_cookies;
@@ -90,17 +90,18 @@ mod tests {
 
     let cookies = get_test_cookies()?;
 
-    let websocket = create_listen_websocket(CreateListenWebsocketArgs {
+    let mut websocket = create_listen_websocket(CreateListenWebsocketArgs {
       cookies: &cookies,
     }).await?;
 
-    let websocket = ClonableWebsocket::new(websocket);
+    //let websocket = GrokWrappedWebsocket::new(websocket);
+    let mut websocket = GrokWebsocket::new(websocket);
 
     println!("Sending...");
     std::io::stdout().flush()?;
 
     let _result = prompt_websocket_image(PromptWebsocketImageArgs {
-      websocket_wrapped: websocket.clone(),
+      websocket: &mut websocket,
       prompt,
     }).await?;
 
@@ -110,7 +111,7 @@ mod tests {
     println!("Polling.");
 
     let request_data = listen_for_websocket_request_id(ListenForWebsocketRequestIdArgs {
-      websocket: &websocket,
+      websocket: &mut websocket,
       timeout: Duration::from_millis(10_000),
     }).await?;
 

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

@@ -1,6 +1,7 @@
-pub mod clonable_websocket;
+pub mod grok_wrapped_websocket;
 pub mod create_listen_websocket;
 pub mod listen_for_websocket_images;
 pub mod listen_for_websocket_request_id;
 pub mod messages;
-pub mod prompt_websocket_image;
+pub mod prompt_websocket_image;
+pub mod grok_websocket;

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

@@ -1,9 +1,9 @@
 use crate::error::grok_error::GrokError;
-use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
+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: ClonableWebsocket,
+  pub websocket: &'a mut GrokWebsocket,
   pub prompt: &'a str,
 }
 
@@ -11,15 +11,16 @@ pub struct PromptWebsocketImageArgs<'a> {
 pub async fn prompt_websocket_image(args: PromptWebsocketImageArgs<'_>) -> Result<(), GrokError> {
   let message = WebsocketClientMessage::new_image_prompt(args.prompt);
 
-  args.websocket_wrapped.send_serializable(&message).await?;
+  args.websocket.send_serializable(&message).await?;
 
   Ok(())
 }
 
 #[cfg(test)]
 mod tests {
-  use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
   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_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;
@@ -40,13 +41,14 @@ mod tests {
       cookies: &cookies,
     }).await?;
 
-    let websocket = ClonableWebsocket::new(websocket);
+    //let websocket = GrokWrappedWebsocket::new(websocket);
+    let mut websocket = GrokWebsocket::new(websocket);
 
     println!("Sending...");
     std::io::stdout().flush()?;
 
     let result = prompt_websocket_image(PromptWebsocketImageArgs {
-      websocket_wrapped: websocket.clone(),
+      websocket: &mut websocket,
       prompt: "a dog riding a motorcycle",
     }).await?;
 

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

@@ -11,6 +11,10 @@ pub enum ArtcraftError {
   // Service errors
   GrokError(GrokError),
   StorytellerError(StorytellerError),
+  // Lock errors
+  RwLockReadError,
+  RwLockWriteError,
+  MutexLockError,
 }
 
 impl Error for ArtcraftError {}
@@ -23,6 +27,9 @@ 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"),
+      Self::MutexLockError => write!(f, "MutexLockError"),
     }
   }
 }

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,9 @@ 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")),
+      ArtcraftError::MutexLockError => Self::AnyhowError(anyhow!("Mutex lock 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_image_prompt_queue::GrokImagePromptQueue;
 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_image_prompt_queue: State<'_, GrokImagePromptQueue>,
   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_image_prompt_queue,
     &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_image_prompt_queue: &GrokImagePromptQueue,
   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_image_prompt_queue,
     &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_image_prompt_queue: &GrokImagePromptQueue,
   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_image_prompt_queue,
+      ).await;
+    }
     Some(ImageModel::Midjourney) => {
       return handle_midjourney(
         app,

crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/generic/handle_image_artcraft.rs

@@ -58,6 +58,7 @@ pub async fn handle_image_artcraft(
     Some(
       ImageModel::Gemini25Flash |
       ImageModel::GptImage1 |
+      ImageModel::GrokImage |
       ImageModel::Midjourney
     ) => {
       return Err(GenerateError::AnyhowError(anyhow!("wrong logic: another branch should handle this: {:?}", request.model)));

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

@@ -0,0 +1,138 @@
+use crate::core::commands::enqueue::generate_error::GenerateError;
+use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
+use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::EnqueueTextToImageRequest;
+use crate::core::events::generation_events::common::GenerationModel;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
+use crate::services::grok::state::grok_image_prompt_queue::{GrokImagePromptQueue, PromptItem};
+use enums::common::generation_provider::GenerationProvider;
+use enums::tauri::tasks::task_type::TaskType;
+use idempotency::uuid::generate_random_uuid;
+use tauri::AppHandle;
+
+pub async fn handle_grok(
+  app: &AppHandle,
+  request: &EnqueueTextToImageRequest,
+  app_env_configs: &AppEnvConfigs,
+  creds_manager: &GrokCredentialManager,
+  grok_image_prompt_queue: &GrokImagePromptQueue,
+) -> 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);
+
+  let job_id = generate_random_uuid();
+
+  let prompt = request.prompt
+      .as_deref()
+      .map(|prompt| prompt.trim().to_string())
+      .unwrap_or_else(|| "".to_string());
+
+  grok_image_prompt_queue.enqueue(PromptItem {
+    task_id: job_id.clone(),
+    prompt,
+  })?;
+
+  Ok(TaskEnqueueSuccess {
+    provider: GenerationProvider::Grok,
+    model: Some(GenerationModel::GrokImage),
+    provider_job_id: Some(job_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,8 @@ 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_image_prompt_queue::GrokImagePromptQueue;
+use crate::services::grok::threads::grok_image_websocket_thread::grok_image_websocket_thread::grok_image_websocket_thread;
 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 +31,7 @@ pub async fn handle_tauri_startup(
   sora_task_queue: SoraTaskQueue,
   mj_creds_manager: MidjourneyCredentialManager,
   grok_creds_manager: GrokCredentialManager,
+  grok_image_prompt_queue: GrokImagePromptQueue,
 ) -> AnyhowResult<()> {
 
   set_app_log_level(
@@ -77,6 +80,16 @@ pub async fn handle_tauri_startup(
     storyteller_creds_manager.clone(),
   ));
 
+  tauri::async_runtime::spawn(grok_image_websocket_thread(
+    app.clone(),
+    app_env_configs.clone(),
+    root.clone(),
+    task_database.clone(),
+    grok_creds_manager.clone(),
+    grok_image_prompt_queue.clone(),
+    storyteller_creds_manager.clone(),
+  ));
+
   tauri::async_runtime::spawn(midjourney_long_polling_thread(
     app.clone(),
     app_env_configs.clone(),

crates/desktop/artcraft/src/lib.rs

@@ -32,6 +32,7 @@ use crate::core::threads::main_window_thread::main_window_thread::main_window_th
 use crate::services::grok::commands::grok_clear_credentials_command::grok_clear_credentials_command;
 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_image_prompt_queue::GrokImagePromptQueue;
 use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
 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;
@@ -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_prompt_queue = GrokImagePromptQueue::new();
+  let grok_prompt_queue_2 = grok_prompt_queue.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_prompt_queue_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_prompt_queue)
     .manage(midjourney_creds_manager)
     .manage(sora_creds_manager)
     .manage(sora_task_queue)

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

@@ -0,0 +1,76 @@
+use crate::core::artcraft_error::ArtcraftError;
+use log::error;
+use std::collections::VecDeque;
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone)]
+pub struct GrokImagePromptQueue {
+  pub prompt_queue: Arc<Mutex<VecDeque<PromptItem>>>,
+}
+
+
+#[derive(Clone, Debug)]
+pub struct PromptItem {
+  /// Local database task ID
+  pub task_id: String,
+  
+  /// Text prompt
+  pub prompt: String,
+}
+
+impl GrokImagePromptQueue {
+  pub fn new() -> Self {
+    Self {
+      prompt_queue: Arc::new(Mutex::new(VecDeque::new())),
+    }
+  }
+  
+  pub fn enqueue(&self, prompt_item: PromptItem) -> Result<(), ArtcraftError> {
+    match self.prompt_queue.lock() {
+      Ok(mut queue) => {
+        queue.push_back(prompt_item);
+        Ok(())
+      },
+      Err(err) => {
+        error!("Error locking prompt queue: {:?}", err);
+        Err(ArtcraftError::MutexLockError)
+      },
+    }
+  }
+  
+  pub fn dequeue(&self) -> Result<Option<PromptItem>, ArtcraftError> {
+    match self.prompt_queue.lock() {
+      Ok(mut queue) => {
+        Ok(queue.pop_front())
+      },
+      Err(err) => {
+        error!("Error locking prompt queue: {:?}", err);
+        Err(ArtcraftError::MutexLockError)
+      },
+    }
+  }
+
+  pub fn is_empty(&self) -> Result<bool, ArtcraftError> {
+    match self.prompt_queue.lock() {
+      Ok(mut queue) => {
+        Ok(queue.is_empty())
+      },
+      Err(err) => {
+        error!("Error locking prompt queue: {:?}", err);
+        Err(ArtcraftError::MutexLockError)
+      },
+    }
+  }
+
+  pub fn len(&self) -> Result<usize, ArtcraftError> {
+    match self.prompt_queue.lock() {
+      Ok(mut queue) => {
+        Ok(queue.len())
+      },
+      Err(err) => {
+        error!("Error locking prompt queue: {:?}", err);
+        Err(ArtcraftError::MutexLockError)
+      },
+    }
+  }
+}

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

@@ -0,0 +1,62 @@
+use crate::core::artcraft_error::ArtcraftError;
+use errors::AnyhowResult;
+use grok_client::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket;
+use log::error;
+use std::sync::{Arc, 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 GrokWebsocketManager {
+  //websocket: Arc<RwLock<Option<GrokWebsocket>>>,
+  websocket: Arc<RwLock<Option<()>>>,
+}
+
+impl GrokWebsocketManager {
+  pub fn new() -> Self {
+    Self {
+      websocket: Arc::new(RwLock::new(None)),
+    }
+  }
+
+  pub fn set_websocket(&self, websocket: GrokWrappedWebsocket) -> Result<(), ArtcraftError> {
+    match self.websocket.write() {
+      Ok(mut guard) => {
+        //*guard = Some(websocket);
+        *guard = Some(());
+        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<GrokWrappedWebsocket>, ArtcraftError> {
+    match self.websocket.read() {
+      Ok(guard) => {
+        //Ok(guard.clone())
+        Ok(None)
+      }
+      Err(err) => {
+        error!("Error reading locked websocket: {}", err);
+        Err(ArtcraftError::RwLockReadError)
+      }
+    }
+  }
+}

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

@@ -1,3 +1,7 @@
 pub mod grok_credential_manager;
+pub mod grok_image_prompt_queue;
 pub mod grok_serializable_state;
-mod grok_credential_holder;
+pub mod grok_websocket_manager;
+
+// Private
+mod grok_credential_holder;

crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs

@@ -0,0 +1,364 @@
+use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
+use crate::core::events::functional_events::text_to_image_generation_complete_event::{GeneratedImage, TextToImageGenerationCompleteEvent};
+use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
+use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+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::task_database::TaskDatabase;
+use crate::core::utils::download_url_to_temp_dir::download_url_to_temp_dir;
+use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
+use crate::services::grok::state::grok_credential_manager::GrokCredentialManager;
+use crate::services::grok::state::grok_image_prompt_queue::{GrokImagePromptQueue, PromptItem};
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use anyhow::anyhow;
+use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
+use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
+use enums::common::generation_provider::GenerationProvider;
+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::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;
+use grok_client::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ImageResults, ListenForWebsocketImagesArgs};
+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 sqlite_tasks::queries::get_task_by_provider_and_provider_job_id::{get_task_by_provider_and_provider_job_id, GetTaskByProviderAndProviderJobIdArgs};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs};
+use sqlite_tasks::queries::task::Task;
+use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
+use sqlite_tasks::queries::update_successful_task_status_with_metadata_by_provider::{update_successful_task_status_with_metadata_by_provider, UpdateSuccessfulTaskByProviderArgs};
+use std::time::Duration;
+use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
+use storyteller_client::endpoints::media_files::get_media_file::get_media_file;
+use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
+use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs, UploadImageMediaFileSuccessResponse};
+use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
+use storyteller_client::error::api_error::ApiError;
+use storyteller_client::error::storyteller_error::StorytellerError;
+use tauri::AppHandle;
+use tokens::tokens::batch_generations::BatchGenerationToken;
+use tokens::tokens::sqlite::tasks::TaskId;
+
+pub async fn grok_image_websocket_thread(
+  app_handle: AppHandle,
+  app_env_configs: AppEnvConfigs,
+  app_data_root: AppDataRoot,
+  task_database: TaskDatabase,
+  creds: GrokCredentialManager,
+  prompt_queue: GrokImagePromptQueue,
+  storyteller_creds_manager: StorytellerCredentialManager,
+) -> ! {
+  loop {
+    let res = inner_loop(
+      &app_handle,
+      &app_env_configs,
+      &app_data_root,
+      &task_database,
+      &creds,
+      &prompt_queue,
+      &storyteller_creds_manager,
+    ).await;
+    if let Err(err) = res {
+      error!("An error occurred: {:?}", err);
+    }
+    // NB: Only sleep if an error occurs.
+    tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
+  }
+}
+
+async fn inner_loop(
+  app_handle: &AppHandle,
+  app_env_configs: &AppEnvConfigs,
+  app_data_root: &AppDataRoot,
+  task_database: &TaskDatabase,
+  grok_creds: &GrokCredentialManager,
+  prompt_queue: &GrokImagePromptQueue,
+  storyteller_creds_manager: &StorytellerCredentialManager,
+) -> AnyhowResult<()> {
+
+  loop {
+    if !grok_creds.do_task_polling()? {
+      tokio::time::sleep(std::time::Duration::from_millis(10_000)).await;
+      continue;
+    }
+
+    let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?;
+
+    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,
+    }).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 {
+        Some(prompt_item) => prompt_item,
+        None => {
+          tokio::time::sleep(std::time::Duration::from_millis(1_000)).await;
+          continue;
+        }
+      };
+
+      info!("Prompting Grok websocket: {}", prompt_item.prompt);
+
+      let _result = prompt_websocket_image(PromptWebsocketImageArgs {
+        websocket: &mut websocket,
+        prompt: &prompt_item.prompt,
+      }).await?;
+
+      let images = listen_for_websocket_images(ListenForWebsocketImagesArgs {
+        websocket: &mut websocket,
+        timeout: Duration::from_millis(30_000),
+      }).await?;
+
+      upload_images_to_storyteller(
+        &app_handle,
+        &app_env_configs,
+        &app_data_root,
+        &task_database,
+        &storyteller_creds_manager,
+        &prompt_item,
+        &images,
+      ).await?
+    }
+
+  }
+}
+
+//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,
+  app_data_root: &AppDataRoot,
+  task_database: &TaskDatabase,
+  storyteller_creds_manager: &StorytellerCredentialManager,
+  prompt_item: &PromptItem,
+  images: &ImageResults,
+) -> AnyhowResult<()> {
+
+  loop {
+    let storyteller_creds = match storyteller_creds_manager.get_credentials()? {
+      Some(creds) => creds,
+      None => {
+        error!("No Storyteller credentials found. Cannot proceed with Grok polling.");
+        tokio::time::sleep(std::time::Duration::from_millis(5_000)).await;
+        continue;
+      }
+    };
+
+    let prompt = prompt_item.prompt.trim().to_string();
+
+    let request = CreatePromptRequest {
+      uuid_idempotency_token: generate_random_uuid(),
+      positive_prompt: Some(prompt),
+      negative_prompt: None,
+      model_type: Some(ModelType::GrokImage),
+      generation_provider: Some(GenerationProvider::Grok),
+    };
+
+    let prompt_response = create_prompt(
+      &app_env_configs.storyteller_host,
+      Some(&storyteller_creds),
+      request
+    ).await?;
+
+    info!("Created prompt: {:?}", &prompt_response.prompt_token);
+
+    // TODO: Move this from clientside to the backend.
+    //  The first upload should produce a batch token that we can reuse.
+    let batch_token = BatchGenerationToken::generate();
+
+    let mut maybe_primary_media_file_token = None;
+
+    for (i, image) in images.images.iter().enumerate() {
+      loop {
+        let url = image.url.to_string();
+        let file = download_url_to_temp_dir(&url, app_data_root).await?;
+
+        info!("Uploading image {} of {} ...", i, images.images.len());
+
+        let result = upload_image_media_file_from_file(UploadImageFromFileArgs {
+          api_host: &app_env_configs.storyteller_host,
+          maybe_creds: Some(&storyteller_creds),
+          path: file.path(),
+          is_intermediate_system_file: false,
+          maybe_prompt_token: Some(&prompt_response.prompt_token),
+          maybe_batch_token: Some(&batch_token),
+        }).await;
+
+        match result {
+          Ok(uploaded) => {
+            if maybe_primary_media_file_token.is_none() {
+              maybe_primary_media_file_token = Some(uploaded.media_file_token);
+            }
+            break; // Break retry loop.
+          }
+          Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => {
+            tokio::time::sleep(std::time::Duration::from_millis(10_000)).await;
+            continue;
+          }
+          Err(err) => {
+            error!("Error uploading image: {}", err);
+            break;
+          }
+        }
+      } // Retry on 429s, etc.
+    }
+
+    let event = GenerationCompleteEvent {
+      //media_file_token: result.media_file_token,
+      action: Some(GenerationAction::GenerateImage),
+      service: GenerationServiceProvider::Grok,
+      model: None,
+    };
+
+    event.send_infallible(&app_handle);
+
+    let task = get_task_by_provider_and_provider_job_id(GetTaskByProviderAndProviderJobIdArgs {
+      db: task_database.get_connection(),
+      provider: GenerationProvider::Grok,
+      provider_job_id: &prompt_item.task_id,
+    }).await?;
+
+    let task = match task {
+      Some(task) => task,
+      None => {
+        warn!("Couldn't find local sqlite task by id: {}", &prompt_item.task_id);
+        return Ok(());
+      }
+    };
+
+    let mut maybe_cdn_url = None;
+    let mut maybe_thumbnail_url_template = None;
+
+    if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() {
+      info!("Looking up file to grab CDN and thumbnail URLs: {:?} ...", media_file_token);
+
+      let lookup_result = get_media_file(
+        &app_env_configs.storyteller_host,
+        media_file_token,
+      ).await;
+      match lookup_result {
+        Ok(response) => {
+          maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string());
+          maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links)
+              .map(|s| s.to_string());
+        }
+        Err(err) => {
+          error!("Failed to look up media file after upload: {:?} (failing open)", err);
+        }
+      }
+    }
+
+    let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
+      db: task_database.get_connection(),
+      task_id: &task.id,
+      maybe_batch_token: Some(&batch_token),
+      maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(),
+      maybe_primary_media_file_class: Some(TaskMediaFileClass::Image),
+      maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(),
+      maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(),
+    }).await?;
+
+    if !updated {
+      return Ok(()); // If anything breaks with queries, don't spam events.
+    }
+
+    send_frontend_ui_update(
+      app_handle,
+      app_env_configs,
+      Some(&storyteller_creds),
+      &task,
+      &batch_token,
+    ).await?;
+
+    return Ok(())
+  }
+}
+
+async fn send_frontend_ui_update(
+  app: &AppHandle,
+  app_env_configs: &AppEnvConfigs,
+  maybe_creds: Option<&StorytellerCredentialSet>,
+  task: &Task,
+  batch_token: &BatchGenerationToken,
+) -> AnyhowResult<()> {
+
+  let result = list_batch_generated_redux_media_files(
+    &app_env_configs.storyteller_host,
+    maybe_creds,
+    batch_token,
+  ).await?;
+
+  if result.media_files.is_empty() {
+    return Err(anyhow!("No media files found for batch token: {}", batch_token));
+  }
+
+  let media_files = result.media_files
+      .into_iter()
+      .map(|file| GeneratedImage {
+        media_token: file.token,
+        cdn_url: file.media_links.cdn_url,
+        maybe_thumbnail_template: file.media_links.maybe_thumbnail_template,
+      })
+      .collect();
+
+  let event = TextToImageGenerationCompleteEvent {
+    generated_images: media_files,
+    maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(),
+    maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(),
+  };
+
+  event.send_infallible(app);
+
+  Ok(())
+}

crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs

@@ -0,0 +1 @@
+pub mod grok_image_websocket_thread;

crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs

@@ -33,7 +33,8 @@ use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest
 use midjourney_client::utils::get_image_url::get_image_url;
 use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
 use once_cell::sync::Lazy;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::{HashMap, HashSet};

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

@@ -1 +1,2 @@
-pub mod grok_video_task_polling;
+pub mod grok_video_task_polling;
+pub mod grok_image_websocket_thread;

crates/desktop/artcraft/src/services/midjourney/threads/events/maybe_handle_text_to_image_complete_event.rs

@@ -7,7 +7,8 @@ use enums::tauri::tasks::task_type::TaskType;
 use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
 use errors::AnyhowResult;
 use log::{error, warn};
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
 use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
 use tauri::AppHandle;

crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs

@@ -27,7 +27,8 @@ use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest
 use midjourney_client::utils::get_image_url::get_image_url;
 use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
 use once_cell::sync::Lazy;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::{HashMap, HashSet};
@@ -362,7 +363,7 @@ async fn upload_midjourney_batch(
     task_id: &local_task.id,
     maybe_batch_token: Some(&batch_token),
     maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(),
-    maybe_primary_media_file_class: Some(TaskMediaFileClass::Image), // TODO: Just images for now.
+    maybe_primary_media_file_class: Some(TaskMediaFileClass::Image),
     maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(),
     maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(),
   }).await?;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_failed_generations.rs

@@ -11,7 +11,7 @@ use errors::AnyhowResult;
 use log::info;
 use openai_sora_client::requests::common::task_id::TaskId;
 use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::PartialTaskResponse;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::Task;
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::HashMap;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs

@@ -28,7 +28,8 @@ use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew
 use openai_sora_client::requests::common::task_id::TaskId;
 use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::{PartialGeneration, PartialTaskResponse};
 use reqwest::Url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::{HashMap, HashSet};

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs

@@ -10,6 +10,7 @@ use crate::core::state::task_database::TaskDatabase;
 use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
 use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 use crate::services::sora::threads::sora_task_polling::helpers::handle_failed_generations::{handle_classic_failed_generations, FailedGeneration};
 use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::{handle_classic_successful_generations, GenerationItem, SuccessfulGeneration};
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
@@ -26,7 +27,8 @@ use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew
 use openai_sora_client::requests::common::task_id::TaskId;
 use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus;
 use reqwest::Url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::{HashMap, HashSet};
 use std::fs::File;
@@ -36,7 +38,6 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
-use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 
 pub async fn poll_classic_sora_tasks(
   app_handle: &AppHandle,

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs

@@ -10,6 +10,7 @@ use crate::core::state::task_database::TaskDatabase;
 use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
 use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 use crate::services::sora::threads::sora_task_polling::helpers::handle_failed_generations::{handle_classic_failed_generations, FailedGeneration};
 use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::{handle_classic_successful_generations, GenerationItem, SuccessfulGeneration};
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
@@ -28,7 +29,8 @@ use openai_sora_client::requests::common::task_id::TaskId;
 use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus;
 use openai_sora_client::requests::list_sora2_drafts::list_sora2_drafts::Draft;
 use reqwest::Url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::{HashMap, HashSet};
 use std::fs::File;
@@ -38,7 +40,6 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
-use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension;
 
 pub async fn poll_sora_2_tasks(
   app_handle: &AppHandle,

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/sora_task_polling_thread.rs

@@ -12,6 +12,8 @@ use crate::services::sora::state::sora_credential_manager::SoraCredentialManager
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
 use crate::services::sora::threads::sora_task_polling::helpers::handle_failed_generations::handle_classic_failed_generations;
 use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::handle_classic_successful_generations;
+use crate::services::sora::threads::sora_task_polling::helpers::poll_classic_sora_tasks::poll_classic_sora_tasks;
+use crate::services::sora::threads::sora_task_polling::helpers::poll_sora_2_tasks::poll_sora_2_tasks;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
 use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
 use enums::common::generation_provider::GenerationProvider;
@@ -26,7 +28,8 @@ use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew
 use openai_sora_client::requests::common::task_id::TaskId;
 use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus;
 use reqwest::Url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::{HashMap, HashSet};
 use std::fs::File;
@@ -36,8 +39,6 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
-use crate::services::sora::threads::sora_task_polling::helpers::poll_classic_sora_tasks::poll_classic_sora_tasks;
-use crate::services::sora::threads::sora_task_polling::helpers::poll_sora_2_tasks::poll_sora_2_tasks;
 
 pub async fn sora_task_polling_thread(
   app_handle: AppHandle,

crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_inpainting_complete_event.rs

@@ -9,7 +9,8 @@ use enums::tauri::tasks::task_type::TaskType;
 use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
 use errors::AnyhowResult;
 use log::{error, warn};
-use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
+use sqlite_tasks::queries::task::Task;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
 use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
 use tauri::AppHandle;

crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_text_to_image_complete_event.rs

@@ -7,7 +7,8 @@ use enums::tauri::tasks::task_type::TaskType;
 use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
 use errors::AnyhowResult;
 use log::{error, warn};
-use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
+use sqlite_tasks::queries::task::Task;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
 use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
 use tauri::AppHandle;

crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_send_background_removal_complete_event.rs

@@ -6,7 +6,8 @@ use enums::tauri::tasks::task_type::TaskType;
 use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
 use errors::AnyhowResult;
 use log::error;
-use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
+use sqlite_tasks::queries::task::Task;
 use tokens::tokens::media_files::MediaFileToken;
 
 pub async fn maybe_send_background_removal_complete_event(

crates/desktop/artcraft/src/services/storyteller/threads/storyteller_activity_thread.rs

@@ -17,15 +17,15 @@ use errors::AnyhowResult;
 use log::{debug, error, info};
 use os_info::Info;
 use reqwest::Url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::time::Instant;
-use storyteller_client::endpoints::analytics::log_active_user::log_active_user;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
-use storyteller_client::error::api_error::ApiError;
-use storyteller_client::error::storyteller_error::StorytellerError;
+use storyteller_client::endpoints::analytics::log_active_user::log_active_user;
 use storyteller_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States};
 use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
+use storyteller_client::error::api_error::ApiError;
+use storyteller_client::error::storyteller_error::StorytellerError;
 use tauri::AppHandle;
 use tokens::tokens::app_session::AppSessionToken;
 

crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs

@@ -27,7 +27,8 @@ use enums::tauri::tasks::task_type::TaskType;
 use errors::AnyhowResult;
 use log::{error, info};
 use reqwest::Url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
+use sqlite_tasks::queries::task::Task;
 use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
 use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
 use std::collections::HashMap;

crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs

@@ -0,0 +1,78 @@
+use crate::connection::TaskDbConnection;
+use crate::error::SqliteTasksError;
+use crate::queries::task::Task;
+use enums::common::generation_provider::GenerationProvider;
+use enums::tauri::tasks::task_model_type::TaskModelType;
+use enums::tauri::tasks::task_status::TaskStatus;
+use enums::tauri::tasks::task_type::TaskType;
+use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
+use sqlx::Error;
+use tokens::tokens::sqlite::tasks::TaskId;
+
+pub struct GetTaskByProviderAndProviderJobIdArgs<'a> {
+  pub db: &'a TaskDbConnection,
+  pub provider: GenerationProvider,
+  pub provider_job_id: &'a str,
+}
+
+
+pub async fn get_task_by_provider_and_provider_job_id(
+  args: GetTaskByProviderAndProviderJobIdArgs<'_>,
+) -> Result<Option<Task>, SqliteTasksError> {
+
+  // TODO(bt,2025-07-12): Fix this. The sqlx mysql queries never required temporaries
+  let temp_provider = args.provider.to_str();
+  let temp_provider_job_id = args.provider_job_id;
+
+  let query = sqlx::query!(r#"
+    SELECT
+      id,
+      task_status,
+      task_type,
+      model_type,
+      provider,
+      provider_job_id,
+      frontend_caller,
+      frontend_subscriber_id,
+      frontend_subscriber_payload
+    FROM tasks
+    WHERE
+      provider = ?
+      AND provider_job_id = ?
+  "#,
+      temp_provider,
+      temp_provider_job_id,
+  );
+
+  // info!("Query: {:?}", query.sql());
+
+  let result = query.fetch_one(args.db.get_pool()).await;
+
+  let record = match result {
+    Ok(record) => record,
+    Err(err) => return match err {
+      Error::RowNotFound => Ok(None),
+      _ => Err(err.into()),
+    },
+  };
+
+  // NB: Not sure why query can't figure out this isn't nullable
+  let provider = record.provider.as_deref()
+      .unwrap_or_else(|| args.provider.to_str());
+
+  Ok(Some(Task {
+    id: TaskId::new_from_str(&record.id),
+    status: TaskStatus::from_str(&record.task_status)?,
+    task_type: TaskType::from_str(&record.task_type)?,
+    model_type: record.model_type
+        .map(|model| TaskModelType::from_str(&model))
+        .transpose()?,
+    provider: GenerationProvider::from_str(provider)?,
+    provider_job_id: record.provider_job_id,
+    frontend_caller: record.frontend_caller
+        .map(|caller| TauriCommandCaller::from_str(&caller))
+        .transpose()?,
+    frontend_subscriber_id: record.frontend_subscriber_id,
+    frontend_subscriber_payload: record.frontend_subscriber_payload,
+  }))
+}

crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs

@@ -1,5 +1,6 @@
 use crate::connection::TaskDbConnection;
 use crate::error::SqliteTasksError;
+use crate::queries::task::{RawTask, Task};
 use enums::common::generation_provider::GenerationProvider;
 use enums::tauri::tasks::task_model_type::TaskModelType;
 use enums::tauri::tasks::task_status::TaskStatus;
@@ -19,33 +20,6 @@ pub struct TaskList {
   pub tasks: Vec<Task>,
 }
 
-#[derive(Debug, Clone)]
-pub struct Task {
-  pub id: TaskId,
-  pub status: TaskStatus,
-  pub task_type: TaskType,
-  pub model_type: Option<TaskModelType>,
-  pub provider: GenerationProvider,
-  pub provider_job_id: Option<String>,
-  pub frontend_caller: Option<TauriCommandCaller>,
-  pub frontend_subscriber_id: Option<String>,
-  pub frontend_subscriber_payload: Option<String>,
-}
-
-#[derive(Debug)]
-#[derive(sqlx::FromRow)]
-struct RawTask {
-  id: String,
-  task_status: String,
-  task_type: String,
-  model_type: Option<String>,
-  provider: String,
-  provider_job_id: Option<String>,
-  frontend_caller: Option<String>,
-  frontend_subscriber_id: Option<String>,
-  frontend_subscriber_payload: Option<String>,
-}
-
 pub async fn list_tasks_by_provider_and_status(
   args: ListTasksByProviderAndStatusArgs<'_>,
 ) -> Result<TaskList, SqliteTasksError> {

crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs

@@ -1,5 +1,6 @@
 use crate::connection::TaskDbConnection;
 use crate::error::SqliteTasksError;
+use crate::queries::task::{RawTask, Task};
 use enums::common::generation_provider::GenerationProvider;
 use enums::tauri::tasks::task_model_type::TaskModelType;
 use enums::tauri::tasks::task_status::TaskStatus;
@@ -18,33 +19,6 @@ pub struct TaskList {
   pub tasks: Vec<Task>,
 }
 
-#[derive(Debug, Clone)]
-pub struct Task {
-  pub id: TaskId,
-  pub status: TaskStatus,
-  pub task_type: TaskType,
-  pub model_type: Option<TaskModelType>,
-  pub provider: GenerationProvider,
-  pub provider_job_id: Option<String>,
-  pub frontend_caller: Option<TauriCommandCaller>,
-  pub frontend_subscriber_id: Option<String>,
-  pub frontend_subscriber_payload: Option<String>,
-}
-
-#[derive(Debug)]
-#[derive(sqlx::FromRow)]
-struct RawTask {
-  id: String,
-  task_status: String,
-  task_type: String,
-  model_type: Option<String>,
-  provider: String,
-  provider_job_id: Option<String>,
-  frontend_caller: Option<String>,
-  frontend_subscriber_id: Option<String>,
-  frontend_subscriber_payload: Option<String>,
-}
-
 pub async fn list_tasks_by_provider_and_tokens(
   args: ListTasksArgs<'_>,
 ) -> Result<TaskList, SqliteTasksError> {

crates/schema/database/sqlite_tasks/src/queries/mod.rs

@@ -1,7 +1,10 @@
 pub mod create_task;
+pub mod get_task_by_provider_and_provider_job_id;
 pub mod list_tasks_by_provider_and_status;
 pub mod list_tasks_by_provider_and_tokens;
 pub mod list_tasks_for_frontend;
 pub mod mark_task_as_dismissed;
+pub mod task;
 pub mod update_successful_task_status_with_metadata;
+pub mod update_successful_task_status_with_metadata_by_provider;
 pub mod update_task_status;

crates/schema/database/sqlite_tasks/src/queries/task.rs

@@ -0,0 +1,33 @@
+use enums::common::generation_provider::GenerationProvider;
+use enums::tauri::tasks::task_model_type::TaskModelType;
+use enums::tauri::tasks::task_status::TaskStatus;
+use enums::tauri::tasks::task_type::TaskType;
+use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
+use tokens::tokens::sqlite::tasks::TaskId;
+
+#[derive(Debug, Clone)]
+pub struct Task {
+  pub id: TaskId,
+  pub status: TaskStatus,
+  pub task_type: TaskType,
+  pub model_type: Option<TaskModelType>,
+  pub provider: GenerationProvider,
+  pub provider_job_id: Option<String>,
+  pub frontend_caller: Option<TauriCommandCaller>,
+  pub frontend_subscriber_id: Option<String>,
+  pub frontend_subscriber_payload: Option<String>,
+}
+
+#[derive(Debug)]
+#[derive(sqlx::FromRow)]
+pub (crate) struct RawTask {
+  pub (crate) id: String,
+  pub (crate) task_status: String,
+  pub (crate) task_type: String,
+  pub (crate) model_type: Option<String>,
+  pub (crate) provider: String,
+  pub (crate) provider_job_id: Option<String>,
+  pub (crate) frontend_caller: Option<String>,
+  pub (crate) frontend_subscriber_id: Option<String>,
+  pub (crate) frontend_subscriber_payload: Option<String>,
+}

crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata_by_provider.rs

@@ -0,0 +1,64 @@
+use enums::common::generation_provider::GenerationProvider;
+use crate::connection::TaskDbConnection;
+use crate::error::SqliteTasksError;
+use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
+use enums::tauri::tasks::task_status::TaskStatus;
+use tokens::tokens::batch_generations::BatchGenerationToken;
+use tokens::tokens::media_files::MediaFileToken;
+
+const SUCCESSFUL_STATUS : &str = TaskStatus::CompleteSuccess.to_str();
+
+pub struct UpdateSuccessfulTaskByProviderArgs<'a> {
+  pub db: &'a TaskDbConnection,
+  pub provider: GenerationProvider,
+  pub provider_job_id: &'a str,
+  pub maybe_batch_token: Option<&'a BatchGenerationToken>,
+  pub maybe_primary_media_file_token: Option<&'a MediaFileToken>,
+  pub maybe_primary_media_file_class: Option<TaskMediaFileClass>,
+  pub maybe_primary_media_file_cdn_url: Option<&'a str>,
+  pub maybe_primary_media_file_thumbnail_url_template: Option<&'a str>,
+}
+
+/// Returns true if rows were updated.
+pub async fn update_successful_task_status_with_metadata_by_provider(
+  args: UpdateSuccessfulTaskByProviderArgs<'_>,
+) -> Result<bool, SqliteTasksError> {
+
+  // TODO(bt,2025-07-12): Fix this. The sqlx mysql queries never required temporaries
+  let provider_temp = args.provider.to_str();
+  let provider_job_id_temp = args.provider_job_id;
+  let maybe_batch_token_temp = args.maybe_batch_token.map(|t| t.as_str());
+  let maybe_primary_media_token_temp = args.maybe_primary_media_file_token.map(|t| t.as_str());
+  let maybe_primary_media_class = args.maybe_primary_media_file_class.map(|c| c.to_str());
+
+  // TODO(bt,2025-07-15): We can't set a LIMIT without a certain compiler flag for SQLite ?
+  let query = sqlx::query!(r#"
+    UPDATE tasks
+    SET
+      task_status = ?,
+      on_complete_batch_token = ?,
+      on_complete_primary_media_file_token = ?,
+      on_complete_primary_media_file_class = ?,
+      on_complete_primary_media_file_cdn_url = ?,
+      on_complete_primary_media_file_thumbnail_url_template = ?
+    WHERE
+      provider = ?
+      AND provider_job_id = ?
+  "#,
+      SUCCESSFUL_STATUS,
+      maybe_batch_token_temp,
+      maybe_primary_media_token_temp,
+      maybe_primary_media_class,
+      args.maybe_primary_media_file_cdn_url,
+      args.maybe_primary_media_file_thumbnail_url_template,
+      provider_temp,
+      provider_job_id_temp,
+  );
+
+  // info!("query: {:?}", query.sql());
+
+  let res = query.execute(args.db.get_pool()).await?;
+  let rows_updated = res.rows_affected() > 0;
+
+  Ok(rows_updated)
+}

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

frontend/libs/model-list/src/lib/lists/ImageModels.ts

@@ -3,6 +3,21 @@ import { ModelCreator } from "../classes/metadata/ModelCreator.js";
 import { ModelTag } from "../classes/metadata/ModelTag.js";
 
 export const IMAGE_MODELS: ImageModel[] = [
+  new ImageModel({
+    id: "grok_image",
+    tauriId: "grok_image",
+    fullName: "Grok",
+    category: "image",
+    creator: ModelCreator.Grok,
+    selectorName: "Grok",
+    selectorDescription: "Fast af",
+    selectorBadges: ["10 sec."],
+    maxGenerationCount: 6,
+    defaultGenerationCount: 6,
+    canUseImagePrompt: false,
+    maxImagePromptCount: 6,
+    progressBarTime: 10000,
+  }),
   new ImageModel({
     id: "midjourney",
     tauriId: "midjourney",