From c3ba2514474884568bedbd4038fc3351def118fe Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 31 Oct 2025 16:01:21 -0400 Subject: [PATCH] hm, sync/send issues with a mutex/arc? --- .../src/datatypes/api/request_id.rs | 6 + .../listen_for_websocket_images.rs | 2 +- .../listen_for_websocket_request_id.rs | 2 +- .../image_websocket/prompt_websocket_image.rs | 4 +- .../artcraft/src/core/artcraft_error.rs | 5 + .../core/commands/enqueue/generate_error.rs | 5 + .../commands/enqueue/task_enqueue_success.rs | 1 + .../enqueue_text_to_image_command.rs | 24 +++ .../enqueue/text_to_image/grok/handle_grok.rs | 155 ++++++++++++++++++ .../enqueue/text_to_image/grok/mod.rs | 1 + .../commands/enqueue/text_to_image/mod.rs | 3 +- .../core/events/generation_events/common.rs | 2 + .../lifecycle/startup/handle_tauri_startup.rs | 2 + crates/desktop/artcraft/src/lib.rs | 6 + .../grok/state/grok_websocket_manager.rs | 60 ++++++- .../artcraft/src/services/grok/state/mod.rs | 4 +- .../enums/src/tauri/tasks/task_model_type.rs | 11 +- 17 files changed, 282 insertions(+), 11 deletions(-) create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs diff --git a/crates/api_clients/grok_client/src/datatypes/api/request_id.rs b/crates/api_clients/grok_client/src/datatypes/api/request_id.rs index ff0f796166..1e06ece469 100644 --- a/crates/api_clients/grok_client/src/datatypes/api/request_id.rs +++ b/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() + } +} 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 6f4211a621..49a13ad9e4 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 @@ -143,7 +143,7 @@ mod tests { std::io::stdout().flush()?; let _result = prompt_websocket_image(PromptWebsocketImageArgs { - websocket_wrapped: websocket.clone(), + websocket_wrapped: &websocket, prompt, }).await?; 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 84c313783d..18bbf6f192 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 @@ -100,7 +100,7 @@ mod tests { std::io::stdout().flush()?; let _result = prompt_websocket_image(PromptWebsocketImageArgs { - websocket_wrapped: websocket.clone(), + websocket_wrapped: &websocket, prompt, }).await?; 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 1f2d4859d8..c9e56935d3 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: GrokWebsocket, + pub websocket_wrapped: &'a GrokWebsocket, pub prompt: &'a str, } @@ -46,7 +46,7 @@ mod tests { std::io::stdout().flush()?; let result = prompt_websocket_image(PromptWebsocketImageArgs { - websocket_wrapped: websocket.clone(), + websocket_wrapped: &websocket, prompt: "a dog riding a motorcycle", }).await?; diff --git a/crates/desktop/artcraft/src/core/artcraft_error.rs b/crates/desktop/artcraft/src/core/artcraft_error.rs index 8830a77973..293b2b6169 100644 --- a/crates/desktop/artcraft/src/core/artcraft_error.rs +++ b/crates/desktop/artcraft/src/core/artcraft_error.rs @@ -11,6 +11,9 @@ pub enum ArtcraftError { // Service errors GrokError(GrokError), StorytellerError(StorytellerError), + // Lock errors + RwLockReadError, + RwLockWriteError, } impl Error for ArtcraftError {} @@ -23,6 +26,8 @@ impl Display for ArtcraftError { Self::IoError(e) => write!(f, "IoError: {:?}", e), Self::GrokError(e) => write!(f, "GrokError: {:?}", e), Self::StorytellerError(e) => write!(f, "StorytellerError: {:?}", e), + Self::RwLockReadError => write!(f, "RwLockReadError"), + Self::RwLockWriteError => write!(f, "RwLockWriteError"), } } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs index 5f16dd9b5e..7ac79c44a3 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs @@ -1,4 +1,5 @@ use crate::core::artcraft_error::ArtcraftError; +use anyhow::anyhow; use base64::DecodeError; use errors::AnyhowError; use grok_client::error::grok_error::GrokError; @@ -79,6 +80,8 @@ pub enum BillingProvider { #[derive(Debug)] pub enum ProviderFailureReason { GrokError(GrokError), + /// NB: The Grok client doesn't say why certain errors (eg. missing fields) happen, so we synthesize this. + GrokJobEnqueueFailed, //Fal(FalErrorPlus), MidjourneyError(MidjourneyError), /// NB: The midjourney client doesn't categorize all errors, so we have to do so on our end. @@ -148,6 +151,8 @@ impl From 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")), } } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index 2d3297f03e..848dc484e6 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/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), 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 3eb152ea63..aa7403e29e 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 @@ -6,6 +6,7 @@ use crate::core::commands::enqueue::text_to_image::generic::handle_image_artcraf use crate::core::commands::enqueue::text_to_image::generic::handle_image_fal::handle_image_fal; use crate::core::commands::enqueue::text_to_image::gpt_image_1::handle_gpt_image_1::handle_gpt_image_1; use crate::core::commands::enqueue::text_to_image::gpt_image_1::handle_gpt_image_1_sora::handle_gpt_image_1_sora; +use crate::core::commands::enqueue::text_to_image::grok::handle_grok::handle_grok; use crate::core::commands::enqueue::text_to_image::midjourney::handle_midjourney::handle_midjourney; use crate::core::commands::response::failure_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus}; use crate::core::commands::response::shorthand::Response; @@ -19,6 +20,8 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::core::state::provider_priority::{Provider, ProviderPriorityStore}; use crate::core::state::task_database::TaskDatabase; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; use crate::services::sora::state::sora_task_queue::SoraTaskQueue; @@ -48,6 +51,8 @@ pub enum ImageModel { FluxPro11Ultra, #[serde(rename = "gpt_image_1")] GptImage1, + #[serde(rename = "grok_image")] + GrokImage, #[serde(rename = "recraft_3")] Recraft3, @@ -144,6 +149,8 @@ pub async fn enqueue_text_to_image_command( provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, mj_creds_manager: State<'_, MidjourneyCredentialManager>, + grok_creds_manager: State<'_, GrokCredentialManager>, + grok_websocket_manager: State<'_, GrokWebsocketManager>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, sora_creds_manager: State<'_, SoraCredentialManager>, sora_task_queue: State<'_, SoraTaskQueue>, @@ -158,6 +165,8 @@ pub async fn enqueue_text_to_image_command( &provider_priority_store, &task_database, &mj_creds_manager, + &grok_creds_manager, + &grok_websocket_manager, &storyteller_creds_manager, &app_env_configs, &sora_creds_manager, @@ -226,6 +235,8 @@ pub async fn handle_request( provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, mj_creds_manager: &MidjourneyCredentialManager, + grok_creds_manager: &GrokCredentialManager, + grok_websocket_manager: &GrokWebsocketManager, storyteller_creds_manager: &StorytellerCredentialManager, app_env_configs: &AppEnvConfigs, sora_creds_manager: &SoraCredentialManager, @@ -240,6 +251,8 @@ pub async fn handle_request( &storyteller_creds_manager, &app_env_configs, &mj_creds_manager, + &grok_creds_manager, + &grok_websocket_manager, &sora_creds_manager, &sora_task_queue, ).await; @@ -274,6 +287,8 @@ pub async fn dispatch_request( storyteller_creds_manager: &StorytellerCredentialManager, app_env_configs: &AppEnvConfigs, mj_creds_manager: &MidjourneyCredentialManager, + grok_creds_manager: &GrokCredentialManager, + grok_websocket_manager: &GrokWebsocketManager, sora_creds_manager: &SoraCredentialManager, sora_task_queue: &SoraTaskQueue, ) -> Result { @@ -306,6 +321,15 @@ pub async fn dispatch_request( sora_task_queue, ).await; } + Some(ImageModel::GrokImage) => { + return handle_grok( + app, + request, + app_env_configs, + grok_creds_manager, + grok_websocket_manager, + ) + } Some(ImageModel::Midjourney) => { return handle_midjourney( app, 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 new file mode 100644 index 0000000000..d82db37b9b --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs @@ -0,0 +1,155 @@ +use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason, ProviderFailureReason}; +use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize}; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent; +use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent; +use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::provider_priority::ProviderPriorityStore; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager; +use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use artcraft_api_defs::generate::image::generate_gpt_image_1_text_to_image::{GenerateGptImage1TextToImageImageQuality, GenerateGptImage1TextToImageImageSize, GenerateGptImage1TextToImageNumImages, GenerateGptImage1TextToImageRequest}; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_type::TaskType; +use grok_client::error::grok_error::GrokError; +use grok_client::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; +use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; +use grok_client::requests::image_websocket::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs}; +use grok_client::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; +use idempotency::uuid::generate_random_uuid; +use log::{error, info, warn}; +use midjourney_client::client::midjourney_hostname::MidjourneyHostname; +use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest}; +use midjourney_client::error::midjourney_api_error::MidjourneyApiError; +use midjourney_client::recipes::channel_id::ChannelId; +use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageError, TextToImageRequest}; +use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session; +use openai_sora_client::recipes::simple_image_gen_with_session_auto_renew::{simple_image_gen_with_session_auto_renew, SimpleImageGenAutoRenewRequest}; +use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages}; +use std::time::Duration; +use storyteller_client::endpoints::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image; +use storyteller_client::endpoints::generate::image::generate_gpt_image_1_text_to_image::generate_gpt_image_1_text_to_image; +use tauri::AppHandle; +use tokens::tokens::media_files::MediaFileToken; + +pub async fn handle_grok( + app: &AppHandle, + request: &EnqueueTextToImageRequest, + app_env_configs: &AppEnvConfigs, + creds_manager: &GrokCredentialManager, + websocket_manager: &GrokWebsocketManager, +) -> Result { + + // TODO: We can request population of the user info if absent or expired. + let websocket = get_websocket( + app, + creds_manager, + websocket_manager, + ).await?; + + let _result = prompt_websocket_image(PromptWebsocketImageArgs { + websocket_wrapped: &websocket, + prompt: request.prompt.as_deref().unwrap_or_else(||""), + }).await?; + + let request_id = listen_for_websocket_request_id(ListenForWebsocketRequestIdArgs { + websocket: &websocket, + timeout: Duration::from_millis(10_000), + }).await?; + + let request_id = match request_id.request_id { + Some(request_id) => request_id.to_string(), + None => { + warn!("No request id from Grok image generation (websocket)"); + return Err(GenerateError::ProviderFailure(ProviderFailureReason::GrokJobEnqueueFailed)); + } + }; + + info!("Successfully enqueued MidJourney. Job token: {}", request_id); + + Ok(TaskEnqueueSuccess { + provider: GenerationProvider::Grok, + model: Some(GenerationModel::GrokImage), + provider_job_id: Some(request_id), + task_type: TaskType::ImageGeneration, + }) +} + +// TODO: +// fn handle_midjourney_errors( +// app: &AppHandle, +// maybe_errors: Option> +// ) -> Result { +// if let Some(errors) = maybe_errors { +// if !errors.is_empty() { +// let messages: Vec = 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 { + 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) +} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs new file mode 100644 index 0000000000..181e756f1a --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs @@ -0,0 +1 @@ +pub (super) mod handle_grok; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs index 073f828a2a..7f981744f8 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs +++ b/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; \ No newline at end of file +mod midjourney; +mod grok; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/events/generation_events/common.rs b/crates/desktop/artcraft/src/core/events/generation_events/common.rs index 688164d404..1ee08c3772 100644 --- a/crates/desktop/artcraft/src/core/events/generation_events/common.rs +++ b/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, 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 967c3c0b60..5d68d6d878 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 @@ -10,6 +10,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::artcraft_platform_info::ArtcraftPlatformInfo; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager; use crate::services::grok::threads::grok_video_task_polling::grok_video_task_polling_thread::grok_video_task_polling_thread; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread; @@ -29,6 +30,7 @@ pub async fn handle_tauri_startup( sora_task_queue: SoraTaskQueue, mj_creds_manager: MidjourneyCredentialManager, grok_creds_manager: GrokCredentialManager, + grok_websocket_manager: GrokWebsocketManager, ) -> AnyhowResult<()> { set_app_log_level( diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 7635aac599..1cd7a5ef4d 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -33,6 +33,7 @@ use crate::services::grok::commands::grok_clear_credentials_command::grok_clear_ use crate::services::grok::commands::grok_get_credential_info_command::grok_get_credential_info_command; use crate::services::grok::commands::grok_open_login_command::grok_open_login_command; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager; use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command; use crate::services::midjourney::commands::midjourney_get_credential_info_command::midjourney_get_credential_info_command; use crate::services::midjourney::commands::midjourney_open_login_command::midjourney_open_login_command; @@ -104,6 +105,9 @@ pub fn run() { let grok_creds_manager = GrokCredentialManager::initialize_from_disk_infallible(&app_data_root); let grok_creds_manager_2 = grok_creds_manager.clone(); + let grok_websocket_manager = GrokWebsocketManager::new(); + let grok_websocket_manager_2= grok_websocket_manager.clone(); + println!("Initializing backend runtime..."); let builder = tauri::Builder::default() @@ -142,6 +146,7 @@ pub fn run() { sora_tasks, midjourney_creds_manager_2, grok_creds_manager_2, + grok_websocket_manager_2, ).await; if let Err(err) = result { @@ -157,6 +162,7 @@ pub fn run() { .manage(app_preferences) .manage(artcraft_platform_info) .manage(grok_creds_manager) + .manage(grok_websocket_manager) .manage(midjourney_creds_manager) .manage(sora_creds_manager) .manage(sora_task_queue) 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 d912d9495c..25934ad2a7 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,7 +1,59 @@ +use crate::core::artcraft_error::ArtcraftError; +use errors::AnyhowResult; use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; -use std::sync::{Arc, RwLock}; +use log::error; +use std::sync::{Arc, LockResult, RwLock}; + +/// NB: This is inefficient because the websockets are locked at two layers. +/// Should be fine for our performance needs, though. #[derive(Clone)] -pub struct GrokWebsocketHolder { - websocket: Arc>, -} \ No newline at end of file +pub struct GrokWebsocketManager { + websocket: Arc>>, +} + +impl GrokWebsocketManager { + pub fn new() -> Self { + Self { + websocket: Arc::new(RwLock::new(None)), + } + } + + pub fn set_websocket(&self, websocket: GrokWebsocket) -> Result<(), ArtcraftError> { + match self.websocket.write() { + Ok(mut guard) => { + *guard = Some(websocket); + Ok(()) + } + Err(err) => { + error!("Error writing locked websocket: {}", err); + Err(ArtcraftError::RwLockWriteError) + } + } + } + + pub fn clear_websocket(&self) -> Result<(), ArtcraftError> { + match self.websocket.write() { + Ok(mut guard) => { + *guard = None; + Ok(()) + } + Err(err) => { + error!("Error writing locked websocket: {}", err); + Err(ArtcraftError::RwLockWriteError) + } + } + } + + pub fn grab_websocket(&self) -> Result, ArtcraftError> { + match self.websocket.read() { + Ok(guard) => { + Ok(guard.clone()) + } + Err(err) => { + error!("Error reading locked websocket: {}", err); + Err(ArtcraftError::RwLockReadError) + } + } + } +} diff --git a/crates/desktop/artcraft/src/services/grok/state/mod.rs b/crates/desktop/artcraft/src/services/grok/state/mod.rs index d2c195dc4d..f467802043 100644 --- a/crates/desktop/artcraft/src/services/grok/state/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/state/mod.rs @@ -1,4 +1,6 @@ pub mod grok_credential_manager; pub mod grok_serializable_state; +pub mod grok_websocket_manager; + +// Private mod grok_credential_holder; -mod grok_websocket_manager; \ No newline at end of file diff --git a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs index 93b00b6c9b..99f08a361f 100644 --- a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs +++ b/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