websockets refactor wip

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

@@ -11,11 +11,11 @@ use wreq::ws::message::Message;
 use wreq::ws::WebSocket;
 
 #[derive(Clone)]
-pub struct ClonableWebsocket {
+pub struct GrokWebsocket {
   pub(crate) websocket: Arc<RwLock<WebSocket>>,
 }
 
-impl ClonableWebsocket {
+impl GrokWebsocket {
   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,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::clonable_websocket::ClonableWebsocket;
+use crate::requests::image_websocket::grok_websocket::GrokWebsocket;
 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 +12,7 @@ use wreq::ws::message::Message;
 const DEFAULT_IMAGE_COUNT: usize = 6;
 
 pub struct ListenForWebsocketImagesArgs<'a> {
-  pub websocket: &'a ClonableWebsocket,
+  pub websocket: &'a GrokWebsocket,
   pub timeout: Duration,
 }
 
@@ -113,7 +113,7 @@ 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::grok_websocket::GrokWebsocket;
   use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs};
   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};
@@ -137,7 +137,7 @@ mod tests {
       cookies: &cookies,
     }).await?;
 
-    let websocket = ClonableWebsocket::new(websocket);
+    let websocket = GrokWebsocket::new(websocket);
 
     println!("Sending...");
     std::io::stdout().flush()?;

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 GrokWebsocket,
   pub timeout: Duration,
 }
 
@@ -70,7 +70,7 @@ 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::grok_websocket::GrokWebsocket;
   use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs};
   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};
@@ -94,7 +94,7 @@ mod tests {
       cookies: &cookies,
     }).await?;
 
-    let websocket = ClonableWebsocket::new(websocket);
+    let websocket = GrokWebsocket::new(websocket);
 
     println!("Sending...");
     std::io::stdout().flush()?;

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

@@ -1,4 +1,4 @@
-pub mod clonable_websocket;
+pub mod grok_websocket;
 pub mod create_listen_websocket;
 pub mod listen_for_websocket_images;
 pub mod listen_for_websocket_request_id;

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_wrapped: GrokWebsocket,
   pub prompt: &'a str,
 }
 
@@ -18,7 +18,7 @@ pub async fn prompt_websocket_image(args: PromptWebsocketImageArgs<'_>) -> Resul
 
 #[cfg(test)]
 mod tests {
-  use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket;
+  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::messages::websocket_server_message::WebsocketServerMessage;
   use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs};
@@ -40,7 +40,7 @@ mod tests {
       cookies: &cookies,
     }).await?;
 
-    let websocket = ClonableWebsocket::new(websocket);
+    let websocket = GrokWebsocket::new(websocket);
 
     println!("Sending...");
     std::io::stdout().flush()?;

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

@@ -0,0 +1,7 @@
+use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket;
+use std::sync::{Arc, RwLock};
+
+#[derive(Clone)]
+pub struct GrokWebsocketHolder {
+  websocket: Arc<RwLock<GrokWebsocket>>,
+}

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

@@ -1,3 +1,4 @@
 pub mod grok_credential_manager;
 pub mod grok_serializable_state;
-mod grok_credential_holder;
+mod grok_credential_holder;
+mod grok_websocket_manager;