refactor so this should work

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

@@ -10,34 +10,28 @@ use std::time::Duration;
 use wreq::ws::message::Message;
 use wreq::ws::WebSocket;
 
-#[derive(Clone)]
 pub struct GrokWebsocket {
-  pub(crate) websocket: Arc<RwLock<WebSocket>>,
+  pub(crate) websocket: WebSocket,
 }
 
 impl GrokWebsocket {
   pub fn new(websocket: WebSocket) -> Self {
     Self {
-      websocket: Arc::new(RwLock::new(websocket)),
+      websocket,
     }
   }
 
-  pub async fn send(&self, message: String) -> Result<(), GrokError> {
-    match self.websocket.write() {
-      Err(err) => Err(GrokClientError::WebsocketLockError.into()),
-      Ok(mut websocket) => {
-        let message = Message::text(message);
-        websocket.send(message)
-            .await
-            .map_err(|err| {
-              GrokClientError::WreqClientError(err)
-            })?;
-        Ok(())
-      },
-    }
+  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>(&self, message: T) -> Result<(), GrokError> {
+  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);
@@ -46,54 +40,44 @@ impl GrokWebsocket {
     self.send(message_json).await
   }
 
-  pub async fn get_response_with_timeout(&self, duration: Duration) -> Result<Option<WebsocketServerMessage>, GrokError> {
-    match self.websocket.write() {
-      Err(_) => Err(GrokClientError::WebsocketLockError.into()),
-      Ok(mut websocket) => {
-        let result = tokio::time::timeout(duration, websocket.try_next()).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 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)
-              },
-            }
-          }
+        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(&self, duration: Duration) -> Result<Option<Message>, GrokClientError> {
-    match self.websocket.write() {
-      Err(_) => Err(GrokClientError::WebsocketLockError),
-      Ok(mut websocket) => {
-        let result = tokio::time::timeout(duration, websocket.try_next()).await;
+  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)
-          }
-        }
-      },
+    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)
+      }
     }
   }
 

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

@@ -0,0 +1,124 @@
+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;
+
+#[derive(Clone)]
+pub struct GrokWrappedWebsocket {
+  pub(crate) websocket: Arc<RwLock<WebSocket>>,
+}
+
+impl GrokWrappedWebsocket {
+  pub fn new(websocket: WebSocket) -> Self {
+    Self {
+      websocket: Arc::new(RwLock::new(websocket)),
+    }
+  }
+
+  pub async fn send(&self, message: String) -> Result<(), GrokError> {
+    match self.websocket.write() {
+      Err(err) => Err(GrokClientError::WebsocketLockError.into()),
+      Ok(mut websocket) => {
+        let message = Message::text(message);
+        websocket.send(message)
+            .await
+            .map_err(|err| {
+              GrokClientError::WreqClientError(err)
+            })?;
+        Ok(())
+      },
+    }
+  }
+
+  pub async fn send_serializable<T: Serialize>(&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(&self, duration: Duration) -> Result<Option<WebsocketServerMessage>, GrokError> {
+    match self.websocket.write() {
+      Err(_) => Err(GrokClientError::WebsocketLockError.into()),
+      Ok(mut websocket) => {
+        let result = tokio::time::timeout(duration, 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(&self, duration: Duration) -> Result<Option<Message>, GrokClientError> {
+    match self.websocket.write() {
+      Err(_) => Err(GrokClientError::WebsocketLockError),
+      Ok(mut websocket) => {
+        let result = tokio::time::timeout(duration, 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/listen_for_websocket_images.rs

@@ -2,6 +2,7 @@ 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::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 GrokWebsocket,
+  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::grok_websocket::GrokWebsocket;
   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 = GrokWebsocket::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,
+      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

@@ -7,7 +7,7 @@ use std::time::{Duration, Instant};
 use wreq::ws::message::Message;
 
 pub struct ListenForWebsocketRequestIdArgs<'a> {
-  pub websocket: &'a GrokWebsocket,
+  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::grok_websocket::GrokWebsocket;
   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 = GrokWebsocket::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,
+      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 grok_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

@@ -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: &'a GrokWebsocket,
+  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::grok_websocket::GrokWebsocket;
   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 = GrokWebsocket::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,
+      websocket: &mut websocket,
       prompt: "a dog riding a motorcycle",
     }).await?;
 

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

@@ -22,7 +22,6 @@ 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::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;
@@ -151,7 +150,6 @@ pub async fn enqueue_text_to_image_command(
   task_database: State<'_, TaskDatabase>,
   mj_creds_manager: State<'_, MidjourneyCredentialManager>,
   grok_creds_manager: State<'_, GrokCredentialManager>,
-  grok_websocket_manager: State<'_, GrokWebsocketManager>,
   grok_image_prompt_queue: State<'_, GrokImagePromptQueue>,
   storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
   sora_creds_manager: State<'_, SoraCredentialManager>,
@@ -168,7 +166,6 @@ pub async fn enqueue_text_to_image_command(
     &task_database,
     &mj_creds_manager,
     &grok_creds_manager,
-    &grok_websocket_manager,
     &grok_image_prompt_queue,
     &storyteller_creds_manager,
     &app_env_configs,
@@ -239,7 +236,6 @@ pub async fn handle_request(
   task_database: &TaskDatabase,
   mj_creds_manager: &MidjourneyCredentialManager,
   grok_creds_manager: &GrokCredentialManager,
-  grok_websocket_manager: &GrokWebsocketManager,
   grok_image_prompt_queue: &GrokImagePromptQueue,
   storyteller_creds_manager: &StorytellerCredentialManager,
   app_env_configs: &AppEnvConfigs,
@@ -256,7 +252,6 @@ pub async fn handle_request(
     &app_env_configs,
     &mj_creds_manager,
     &grok_creds_manager,
-    &grok_websocket_manager,
     &grok_image_prompt_queue,
     &sora_creds_manager,
     &sora_task_queue,
@@ -293,7 +288,6 @@ pub async fn dispatch_request(
   app_env_configs: &AppEnvConfigs,
   mj_creds_manager: &MidjourneyCredentialManager,
   grok_creds_manager: &GrokCredentialManager,
-  grok_websocket_manager: &GrokWebsocketManager,
   grok_image_prompt_queue: &GrokImagePromptQueue,
   sora_creds_manager: &SoraCredentialManager,
   sora_task_queue: &SoraTaskQueue,
@@ -333,7 +327,6 @@ pub async fn dispatch_request(
         request,
         app_env_configs,
         grok_creds_manager,
-        grok_websocket_manager,
         grok_image_prompt_queue,
       ).await;
     }

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

@@ -4,7 +4,6 @@ use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command
 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_websocket_manager::GrokWebsocketManager;
 use enums::common::generation_provider::GenerationProvider;
 use enums::tauri::tasks::task_type::TaskType;
 use idempotency::uuid::generate_random_uuid;
@@ -16,7 +15,6 @@ pub async fn handle_grok(
   request: &EnqueueTextToImageRequest,
   app_env_configs: &AppEnvConfigs,
   creds_manager: &GrokCredentialManager,
-  websocket_manager: &GrokWebsocketManager,
   grok_image_prompt_queue: &GrokImagePromptQueue,
 ) -> Result<TaskEnqueueSuccess, GenerateError> {
 
@@ -47,14 +45,14 @@ pub async fn handle_grok(
   //};
   //
   //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(&prompt)?;
 
   Ok(TaskEnqueueSuccess {

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

@@ -11,7 +11,6 @@ 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::state::grok_websocket_manager::GrokWebsocketManager;
 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;
@@ -32,7 +31,6 @@ pub async fn handle_tauri_startup(
   sora_task_queue: SoraTaskQueue,
   mj_creds_manager: MidjourneyCredentialManager,
   grok_creds_manager: GrokCredentialManager,
-  grok_websocket_manager: GrokWebsocketManager,
   grok_image_prompt_queue: GrokImagePromptQueue,
 ) -> AnyhowResult<()> {
 

crates/desktop/artcraft/src/lib.rs

@@ -34,7 +34,6 @@ use crate::services::grok::commands::grok_get_credential_info_command::grok_get_
 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::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;
@@ -106,9 +105,6 @@ 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();
-
   let grok_prompt_queue = GrokImagePromptQueue::new();
   let grok_prompt_queue_2 = grok_prompt_queue.clone();
 
@@ -150,7 +146,6 @@ pub fn run() {
           sora_tasks,
           midjourney_creds_manager_2,
           grok_creds_manager_2,
-          grok_websocket_manager_2,
           grok_prompt_queue_2,
         ).await;
 
@@ -167,7 +162,6 @@ pub fn run() {
     .manage(app_preferences)
     .manage(artcraft_platform_info)
     .manage(grok_creds_manager)
-    .manage(grok_websocket_manager)
     .manage(grok_prompt_queue)
     .manage(midjourney_creds_manager)
     .manage(sora_creds_manager)

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

@@ -1,6 +1,6 @@
 use crate::core::artcraft_error::ArtcraftError;
 use errors::AnyhowResult;
-use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket;
+use grok_client::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket;
 use log::error;
 use std::sync::{Arc, RwLock};
 
@@ -20,7 +20,7 @@ impl GrokWebsocketManager {
     }
   }
 
-  pub fn set_websocket(&self, websocket: GrokWebsocket) -> Result<(), ArtcraftError> {
+  pub fn set_websocket(&self, websocket: GrokWrappedWebsocket) -> Result<(), ArtcraftError> {
     match self.websocket.write() {
       Ok(mut guard) => {
         //*guard = Some(websocket);
@@ -47,7 +47,7 @@ impl GrokWebsocketManager {
     }
   }
 
-  pub fn grab_websocket(&self) -> Result<Option<GrokWebsocket>, ArtcraftError> {
+  pub fn grab_websocket(&self) -> Result<Option<GrokWrappedWebsocket>, ArtcraftError> {
     match self.websocket.read() {
       Ok(guard) => {
         //Ok(guard.clone())

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

@@ -12,7 +12,7 @@ use enums::common::model_type::ModelType;
 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;
@@ -24,6 +24,7 @@ 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 grok_client::requests::image_websocket::grok_websocket::GrokWebsocket;
 use tokens::tokens::batch_generations::BatchGenerationToken;
 
 pub async fn grok_image_websocket_thread(
@@ -83,7 +84,7 @@ async fn inner_loop(
       cookies: &cookie_header,
     }).await?;
 
-    let websocket = GrokWebsocket::new(websocket);
+    let mut websocket = GrokWebsocket::new(websocket);
 
     //poll_task_loop(
     //  &app_handle,
@@ -123,12 +124,12 @@ async fn inner_loop(
       info!("Prompting Grok websocket: {}", prompt);
 
       let _result = prompt_websocket_image(PromptWebsocketImageArgs {
-        websocket_wrapped: &websocket,
+        websocket: &mut websocket,
         prompt: &prompt,
       }).await?;
 
       let images = listen_for_websocket_images(ListenForWebsocketImagesArgs {
-        websocket: &websocket,
+        websocket: &mut websocket,
         timeout: Duration::from_millis(30_000),
       }).await?;