From c41eb800b3bed11ecb7b72148e8f4e565824b9d9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 31 Oct 2025 20:24:52 -0400 Subject: [PATCH] refactor so this should work --- .../image_websocket/grok_websocket.rs | 100 ++++++-------- .../image_websocket/grok_wrapped_websocket.rs | 124 ++++++++++++++++++ .../listen_for_websocket_images.rs | 12 +- .../listen_for_websocket_request_id.rs | 13 +- .../src/requests/image_websocket/mod.rs | 5 +- .../image_websocket/prompt_websocket_image.rs | 12 +- .../enqueue_text_to_image_command.rs | 7 - .../enqueue/text_to_image/grok/handle_grok.rs | 8 +- .../lifecycle/startup/handle_tauri_startup.rs | 2 - crates/desktop/artcraft/src/lib.rs | 6 - .../grok/state/grok_websocket_manager.rs | 6 +- .../grok_image_websocket_thread.rs | 9 +- 12 files changed, 201 insertions(+), 103 deletions(-) create mode 100644 crates/api_clients/grok_client/src/requests/image_websocket/grok_wrapped_websocket.rs diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs b/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs index e607b1678c..9780099bb2 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs +++ b/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>, + 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(&self, message: T) -> Result<(), GrokError> { + pub async fn send_serializable(&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, 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, 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, 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, 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) + } } } diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/grok_wrapped_websocket.rs b/crates/api_clients/grok_client/src/requests/image_websocket/grok_wrapped_websocket.rs new file mode 100644 index 0000000000..712d1878f9 --- /dev/null +++ b/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>, +} + +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(&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, 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, 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 { + // 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, 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), GrokClientError> { + // match self.websocket.read() { + // Err(_) => Err(GrokClientError::WebsocketLockError), + // Ok(websocket) => Ok(websocket.size_hint()) + // } + // } +} diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs index 49a13ad9e4..360107db13 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs +++ b/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?; diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs index 18bbf6f192..a0d0a38999 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs +++ b/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?; diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs b/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs index b15666c033..f6c3676adb 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs +++ b/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; \ No newline at end of file +pub mod prompt_websocket_image; +pub mod grok_websocket; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs b/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs index c9e56935d3..7b223f5e5e 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs +++ b/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?; diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs index c19a7b4f23..87a3915294 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs +++ b/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; } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs index 58dd1ff4d7..b3d7ce80b7 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs +++ b/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 { @@ -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 { diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs index 791b7ed6b1..3c33ef011b 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs +++ b/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<()> { diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 6c4ee490fa..511c694205 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/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) diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs index ff6904a7ff..9e3f5c8e97 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs +++ b/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, ArtcraftError> { + pub fn grab_websocket(&self) -> Result, ArtcraftError> { match self.websocket.read() { Ok(guard) => { //Ok(guard.clone()) diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs index 5c5a7f4f60..e8913a5157 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs +++ b/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?;